From 985ddfa916779a7b1aaf92cf9c142204a65237b6 Mon Sep 17 00:00:00 2001
From: zjgarvey <zjgarvey@gmail.com>
Date: Thu, 31 Oct 2024 14:44:13 -0500
Subject: [PATCH] Add a info cast to prims.squeeze decomposition

---
 lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
index 1fefb59a4cac..1825a151aab9 100644
--- 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();
   }
 };