Fix failing test

llvm · Dec 10, 2024 · 99943ef · 99943ef
1 parent fd771f0
commit 99943ef
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/lib/Conversion/TorchToLinalg/Linear.cpp b/lib/Conversion/TorchToLinalg/Linear.cpp
@@ -1313,6 +1313,15 @@ class ConvertAtenConvolutionOp : public OpConversionPattern<AtenConvolutionOp> {
         conv = torch_to_linalg::convertTensorToElementType(rewriter, loc, conv,
                                                            resultElementType);
       }
+      if (is1DGroupConv) {
+        // Squeezing the last dim of the result of conv.
+        auto squeezeInputInfo = squeezeTensor(rewriter, op, conv, /*dim=*/-1);
+        if (failed(squeezeInputInfo)) {
+          return rewriter.notifyMatchFailure(
+              op, "cannot generate unsqueeze tensor");
+        }
+        conv = squeezeInputInfo.value();
+      }
       rewriter.replaceOpWithNewOp<tensor::CastOp>(op, newResultType, conv);
       return success();
     }