fix sequences_pre

Signed-off-by: Charlene Yang <[email protected]>
NVIDIA · Feb 25, 2025 · f13b861 · f13b861
1 parent 583b76f
commit f13b861
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/transformer_engine/pytorch/inference.py b/transformer_engine/pytorch/inference.py
@@ -320,7 +320,8 @@ def pre_step(
         self.batch_size = len(step_dict)
 
         self.sequences = self.cache_manager.pre_step(step_dict)
-        for k, v in enumerate(self.sequences):
+        self.sequences_pre = OrderedDict()
+        for k, v in self.sequences.items():
             self.sequences_pre[k] = v - self.step_dict[k]
 
         actual_batch_size = len(step_dict)