demo works

pytorch · Aug 8, 2023 · d6506ed · d6506ed
1 parent e87c435
commit d6506ed
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/selective2d/idemo.py b/examples/selective2d/idemo.py
@@ -327,8 +327,8 @@ def pp_and_tp_selective(
                 args.device,
                 pp_groups,
                 example_inputs=[X, Y],
-                #num_stages=num_stages,
-                #schedule="TwoLevel",
+                num_stages=num_stages,
+                schedule="TwoLevel",
         )
     else: 
         output_chunk_spec = (TensorChunkSpec(0), sum_reducer)