From 34b9aa6910da993c94a876a797e64c5c6e6a12ea Mon Sep 17 00:00:00 2001 From: Cheng-Hsin Weng Date: Sun, 9 Aug 2026 20:05:49 +0800 Subject: [PATCH] Qualcomm AI Engine Direct - Change transpose output from NCH1 to NC1H before conv --- backends/qualcomm/utils/utils.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 0eacead775a..77ba97090b6 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -168,11 +168,20 @@ def __init__(self, weight, bias=None): def forward(self, x): rank = x.dim() - x = x.reshape(*x.shape, 1) if rank == 3 else x.reshape(1, *x.shape, 1) - x = torch.transpose(x, 1, 2) + if rank == 3: + bsz, dim0, dim1 = x.size() + x = torch.reshape(x, (bsz, dim0, 1, dim1)) + else: + dim0, dim1 = x.size() + x = torch.reshape(x, (1, dim0, 1, dim1)) + x = torch.transpose(x, 1, 3) res = self.conv(x) - res = torch.transpose(res, 1, 2) - res = res.squeeze(-1) if rank == 3 else res.reshape(*res.shape[1:3]) + res = res.permute(0, 3, 1, 2) + res = ( + res.squeeze(-1) + if rank == 3 + else res.reshape(dim0, self.conv.weight.shape[0]) + ) return res def replace_linear(module: torch.nn.Module):