fix typo

EleutherAI · Oct 7, 2024 · 94d31f6 · 94d31f6
1 parent 774eb58
commit 94d31f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/training.py b/megatron/training.py
@@ -576,7 +576,7 @@ def forward_step(
         return model.eval_batch(data_iterator, return_logits=return_logits)
 
     # Get the batch.
-    if neox_args.memory_profiling and neox_args.it:
+    if neox_args.memory_profiling and neox_args.iteration:
         torch.cuda.nvtx.range_push(f"Get batch")
     if timers is not None:
         timers("batch generator").start()