Add a info cast to prims.squeeze decomposition

llvm · Oct 31, 2024 · 985ddfa · 985ddfa
1 parent a6292f3
commit 985ddfa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -8816,7 +8816,8 @@ class DecomposePrimsSqueezeOp : public OpRewritePattern<PrimsSqueezeOp> {
       }
       result = *squeezeTensorInfo;
     }
-    rewriter.replaceOp(op, result);
+    rewriter.replaceOpWithNewOp<Torch::TensorStaticInfoCastOp>(op, op.getType(),
+                                                               result);
     return success();
   }
 };