right imports

ananthsub · ananthsub · commit 8654eaaeb4bd · 2025-06-04T10:04:31.000-07:00
Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;
diff --git a/nemo_lm/models/__init__.py b/nemo_lm/models/__init__.py
@@ -20,7 +20,7 @@
     TorchFullyShardedDataParallel,
 )
 from megatron.core.enums import ModelType
-from megatron.core.fp8_utils import is_float8tensor
+from megatron.core.fp8_utils import correct_amax_history_if_needed
 from megatron.core.transformer.module import Float16Module, MegatronModule
 
 from nemo_lm.models.gpt import GPTConfig

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@`
`20`	`20`	`TorchFullyShardedDataParallel,`
`21`	`21`	`)`
`22`	`22`	`from megatron.core.enums import ModelType`
`23`		`-from megatron.core.fp8_utils import is_float8tensor`
	`23`	`+from megatron.core.fp8_utils import correct_amax_history_if_needed`
`24`	`24`	`from megatron.core.transformer.module import Float16Module, MegatronModule`
`25`	`25`
`26`	`26`	`from nemo_lm.models.gpt import GPTConfig`