fix mamba config issue

Signed-off-by: Jiaqi Zeng <[email protected]>
NVIDIA · Nov 18, 2024 · c1f1e8b · c1f1e8b
1 parent be8a8d3
commit c1f1e8b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/nlp/gpt/train_gpt_sft.py b/examples/nlp/gpt/train_gpt_sft.py
@@ -105,7 +105,7 @@ def main(cfg) -> None:
         train_data_cfg,
         ptl_model.tokenizer,
         num_samples,
-        is_mamba=cfg.model.mamba_hybrid,
+        is_mamba=cfg.model.get("mamba_hybrid", False),
         answer_only_loss=True,
         is_chat=cfg.model.data.chat,
         special_tokens=cfg.model.data.chat_prompt_tokens,
@@ -118,7 +118,7 @@ def main(cfg) -> None:
         val_data_cfg,
         ptl_model.tokenizer,
         num_samples,
-        is_mamba=cfg.model.mamba_hybrid,
+        is_mamba=cfg.model.get("mamba_hybrid", False),
         answer_only_loss=True,
         is_chat=cfg.model.data.chat,
         special_tokens=cfg.model.data.chat_prompt_tokens,