Fix Gemma2 Attention init args (#11792)

NVIDIA · Jan 9, 2025 · 6db1bfe · 6db1bfe
1 parent 84b2bf0
commit 6db1bfe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py b/nemo/collections/nlp/models/language_modeling/megatron/gemma2/gemma2_modules.py
@@ -67,7 +67,7 @@ def __init__(
         attn_mask_type: AttnMaskType,
         attention_type: str,
         attention_dropout: float = None,
-        cp_comm_type: str = None,
+        **kwargs,
     ):
         super().__init__(config=config)