Set parallelism correctly when using VPP (#10204)

NVIDIA · Aug 20, 2024 · d4f02b5 · d4f02b5
1 parent 54458fa
commit d4f02b5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/lightning/_strategy_lib.py b/nemo/lightning/_strategy_lib.py
@@ -140,6 +140,8 @@ def set_model_parallel_attributes(model, parallelism):
             if not hasattr(config, attr_name):
                 continue
             setattr(config, attr_name, getattr(parallelism, attr_name))
+            if hasattr(config, "__io__"):
+                setattr(config.__io__, attr_name, getattr(parallelism, attr_name))
 
         return config