Fix param input order for cudagraph

Signed-off-by: yifeis-nv <[email protected]>
NVIDIA · Aug 28, 2024 · 4235bfe · 4235bfe
1 parent 7fc50f4
commit 4235bfe
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
@@ -171,8 +171,8 @@ def _make_graphed_callables(
         ]
     else:
         per_callable_module_params = []
-        for c in callables:
-            for i in range(num_microbatches):
+        for i in range(num_microbatches):
+            for c in callables:
                 per_callable_module_params.append(
                     tuple(c.parameters()) if isinstance(c, torch.nn.Module) else ()
                 )