destruct singleton on megatron-strategy dtor

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Oct 30, 2024 · 40113c7 · 40113c7
1 parent a8fd3d6
commit 40113c7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/lightning/pytorch/strategies/megatron_strategy.py b/nemo/lightning/pytorch/strategies/megatron_strategy.py
@@ -609,6 +609,8 @@ def predict_step(self, dataloader_iter, *args: Any, **kwargs: Any) -> STEP_OUTPU
     @override
     def teardown(self) -> None:
         super().teardown()
+        from nemo.utils.metaclasses import Singleton
+        Singleton._Singleton__instances = {}
 
     @override
     def model_sharded_context(self) -> ContextManager: