fix test

Signed-off-by: ashors1 <[email protected]>
NVIDIA · Aug 2, 2024 · b9306e8 · b9306e8
1 parent 33c965e
commit b9306e8
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/lightning/test_dist_ckpt.py b/tests/lightning/test_dist_ckpt.py
@@ -10,9 +10,8 @@
 import nemo.lightning as nl
 from nemo.collections import llm
 from nemo.lightning.io.pl import MegatronCheckpointIO
-from nemo.utils.apex_utils import _reconfigure_microbatch_calculator
 from nemo.utils.callbacks.dist_ckpt_io import AsyncFinalizableCheckpointIO, AsyncFinalizerCallback
-
+from megatron.core.num_microbatches_calculator import reconfigure_num_microbatches_calculator
 
 def _get_strategy():
     strategy = nl.MegatronStrategy(
@@ -41,7 +40,7 @@ def get_model_and_data():
         seq_length=seq_length,
         apply_query_key_layer_scaling=1,
     )
-    _reconfigure_microbatch_calculator(
+    reconfigure_num_microbatches_calculator(
         0,
         None,
         global_batch_size,