keep order

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

nemo/tron/setup.py

@@ -101,6 +101,8 @@ def setup(
     print_rank_0("time to initialize megatron (seconds): {:.3f}".format(time.time() - state.start_time))
     barrier_and_log("after megatron is initialized")
 
+    checkpointing_context = _init_checkpointing_context(cfg)
+
     # Tokenizer
     timers("tokenizer-setup", log_level=0).start(barrier=True)
     tokenizer = build_tokenizer(
@@ -136,8 +138,6 @@ def setup(
     timers("model-and-optimizer-setup").stop()
     barrier_and_log("after model, optimizer, and learning rate scheduler are built")
 
-    checkpointing_context = _init_checkpointing_context(cfg)
-
     # Load checkpoint if applicable
     if (cfg.checkpoint_config.load is not None or cfg.checkpoint_config.pretrained_checkpoint is not None) and (
         checkpoint_exists(cfg.checkpoint_config.load) or checkpoint_exists(cfg.checkpoint_config.pretrained_checkpoint)