update for loop indexing

nod-ai · Aug 7, 2024 · 35209fb · 35209fb
1 parent d0da328
commit 35209fb
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/models/turbine_models/custom_models/stateless_llama.py b/models/turbine_models/custom_models/stateless_llama.py
@@ -273,8 +273,8 @@ def run_forward(self, x=AbstractTensor(1, 1, dtype=torch.int64)):
                     "state_arg0_dim", max=MAX_STEP_SEQ - 1
                 )
                 dynamic_shapes_forw = {"arg0_1": None, "arg1_1": {1: state_arg0_dim}}
-                for dim_number in range(1, len(state_arg)):
-                    current_dim_dict = {f"arg{dim_number + 1}_1": {1: state_arg0_dim}}
+                for state_arg_idx in range(2, len(state_arg) + 1):
+                    current_dim_dict = {f"arg{state_arg_idx}_1": {1: state_arg0_dim}}
                     dynamic_shapes_forw = {**dynamic_shapes_forw, **current_dim_dict}
                 token, *state_update = self.forward(
                     x, *state_arg, dynamic_shapes=dynamic_shapes_forw
@@ -352,8 +352,8 @@ def run_cached_initialize(
                     "arg0_1": {1: x_dim},
                     "arg1_1": {1: state_arg0_dim1},
                 }
-                for dim_number in range(1, len(state_arg)):
-                    current_dim_dict = {f"arg{dim_number + 1}_1": {1: state_arg0_dim1}}
+                for state_arg_idx in range(2, len(state_arg) + 1):
+                    current_dim_dict = {f"arg{state_arg_idx}_1": {1: state_arg0_dim1}}
                     dynamic_shapes_forw = {**dynamic_shapes_forw, **current_dim_dict}
                 token, *state = self.cached_initialize(
                     x, *state_arg, dynamic_shapes=dynamic_shapes_forw