diff --git a/lib/Conversion/TorchToLinalg/DataMovement.cpp b/lib/Conversion/TorchToLinalg/DataMovement.cpp
index 469d82532..f1f1d7329 100644
--- a/lib/Conversion/TorchToLinalg/DataMovement.cpp
+++ b/lib/Conversion/TorchToLinalg/DataMovement.cpp
@@ -1006,7 +1006,9 @@ public:
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
       return failure();
-    rewriter.replaceOp(op, adaptor.self());
+
+    Type resultType = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, adaptor.self());
     return success();
   }
 };
@@ -1074,7 +1076,8 @@ public:
                            })
                        ->getResult(0);
 
-    rewriter.replaceOp(op, result);
+    Type resultType = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, result);
     return success();
   }
 };
diff --git a/lib/Conversion/TorchToLinalg/Uncategorized.cpp b/lib/Conversion/TorchToLinalg/Uncategorized.cpp
index 77d170d55..fcb766515 100644
--- a/lib/Conversion/TorchToLinalg/Uncategorized.cpp
+++ b/lib/Conversion/TorchToLinalg/Uncategorized.cpp
@@ -1538,7 +1538,7 @@ public:
                 })
             .getResult(0);
 
-    rewriter.replaceOp(op, finalRes);
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, finalRes);
     return success();
   }
 };