import fixes

Signed-off-by: Ananth Subramaniam <[email protected]>
ananthsub · Dec 10, 2024 · f313d18 · f313d18
1 parent a0af896
commit f313d18
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/utils/callbacks/dist_ckpt_io.py b/nemo/utils/callbacks/dist_ckpt_io.py
@@ -26,8 +26,8 @@
 from lightning.pytorch import Callback
 from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
 
+from nemo.lightning.ckpt_utils import preprocess_common_state_dict_before_consistency_check
 from nemo.utils import logging
-from nemo.utils.dist_checkpointing import preprocess_common_state_dict_before_consistency_check
 
 try:
     from megatron.core import dist_checkpointing