Remove unused import

NVIDIA · Jun 19, 2024 · 63ad099 · 63ad099
1 parent fd87145
commit 63ad099
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/nemo/lightning/io/pl.py b/nemo/lightning/io/pl.py
@@ -8,7 +8,6 @@
 from lightning_fabric.plugins.io.checkpoint_io import CheckpointIO
 from lightning_fabric.utilities.cloud_io import get_filesystem
 from lightning_fabric.utilities.types import _PATH
-from megatron.core.dist_checkpointing.strategies import tensorstore
 from torch import nn
 from typing_extensions import Self, override
 

diff --git a/nemo/lightning/pytorch/strategies.py b/nemo/lightning/pytorch/strategies.py
@@ -478,7 +478,7 @@ def load_model_state_dict(self, checkpoint: Mapping[str, Any], strict: bool = Tr
 
             mcore_model = self.lightning_module.module
             current = self.model[0]
-            n_nesting = 1
+            n_nesting = 2
             while current != mcore_model:
                 current = current.module
                 n_nesting += 1