diff --git a/csrc/preseg_passes/remove_bcast_squeeze.cpp b/csrc/preseg_passes/remove_bcast_squeeze.cpp
index ce0706dd74c..0dfd782686d 100644
--- a/csrc/preseg_passes/remove_bcast_squeeze.cpp
+++ b/csrc/preseg_passes/remove_bcast_squeeze.cpp
@@ -163,13 +163,13 @@ TensorView* replayAxisOp(
   switch (simple_op_type) {
     case AxisOp::PRESERVE:
       // This is equivalent to a set Op
-      replacement = tv;
+      return tv;
       break;
     case AxisOp::SQUEEZE:
-      replacement = squeeze(tv, nonPreservedDims(axis_ops));
+      return squeeze(tv, nonPreservedDims(axis_ops));
       break;
     case AxisOp::BROADCAST:
-      replacement = broadcast(tv, nonPreservedDims(axis_ops));
+      return broadcast(tv, nonPreservedDims(axis_ops));
       break;
   }
 }