Fix Gemma2 Attention init args (#11792)

NVIDIA · Jan 10, 2025 · d4e61cd · d4e61cd
1 parent 1e6ade2
commit d4e61cd
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)