Merge branch 'main' of https://github.com/NVIDIA/NeMo into athitten/a…

…lways_save_context
NVIDIA · Oct 24, 2024 · 353b2a6 · 353b2a6
2 parents c944e54 + 956406e
commit 353b2a6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/lightning/pytorch/strategies/megatron_strategy.py b/nemo/lightning/pytorch/strategies/megatron_strategy.py
@@ -267,7 +267,7 @@ def __init__(
     def connect(self, model: pl.LightningModule) -> None:
         super().connect(model)
 
-        assert not hasattr(model, 'is_hf_model'), "Cannot use HfAutoModelForCausalLM with MegatronParallel"
+        assert not 'is_hf_model' in model.__dict__, "Cannot use HfAutoModelForCausalLM with MegatronParallel"
 
         _maybe_mcore_config = _strategy_lib.set_model_parallel_attributes(model, self.parallelism)
         if _maybe_mcore_config: