modify loop op

MasterJH5574 · Mar 24, 2024 · fb2d7ba · fb2d7ba
1 parent fcfe005
commit fb2d7ba
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/frontend/control_flow.py b/frontend/control_flow.py
@@ -40,8 +40,7 @@ def forward(self, *values: Any) -> Any:
         loop_carry = values[self.num_read_only_param:]
         while iter_num < self.num_iter:
             # and cond.item():
-            loop_carry = self.body(torch.tensor(iter_num), *read_only,
-                                   *loop_carry)
+            loop_carry = self.body(iter_num, *read_only, *loop_carry)
             # cond, *loop_carry = self.body(iter_num, cond, *read_only,
             #                               *loop_carry)
             iter_num += 1