add TODO in llama forward input shape (#195)

nod-ai · Nov 21, 2023 · d69f658 · d69f658
1 parent c350ba6
commit d69f658
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/turbine_models/custom_models/stateless_llama.py b/python/turbine_models/custom_models/stateless_llama.py
@@ -135,7 +135,7 @@ def run_initialize(self, x=AbstractTensor(BATCH_SIZE, None, dtype=torch.int64)):
                     self.global_state, slice_of_state, i, 0, 0, 0, 0
                 )
             return token
-
+        #TODO: Change None to AbstractTensor(1, 1, dtype=torch.int64) and debug dynamo constraints
         def run_forward(self, x=AbstractTensor(1, None, dtype=torch.int64)):
             state_arg = slice_up_to_step(
                 self.global_state, self.global_seq_step, HEADS, HIDDEN_DIM