Change to kwargs (#5475)

Signed-off-by: MaximumEntropy <[email protected]> Signed-off-by: MaximumEntropy <[email protected]>
NVIDIA · Nov 22, 2022 · 4a523ad · 4a523ad
1 parent 8552c95
commit 4a523ad
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/transformer.py b/nemo/collections/nlp/modules/common/megatron/transformer.py
@@ -2185,13 +2185,13 @@ def custom_forward(*inputs):
                     for index in range(start, end):
                         layer = self._get_layer(index)
                         hidden_states = layer(
-                            hidden_states,
-                            attention_mask,
-                            encoder_output,
-                            enc_dec_attn_mask,
-                            rotary_pos_emb,
-                            self_attention_relative_position_bias,
-                            cross_attention_relative_position_bias,
+                            hidden_states=hidden_states,
+                            attention_mask=attention_mask,
+                            encoder_output=encoder_output,
+                            enc_dec_attn_mask=enc_dec_attn_mask,
+                            rotary_pos_emb=rotary_pos_emb,
+                            self_attention_relative_position_bias=self_attention_relative_position_bias,
+                            cross_attention_relative_position_bias=cross_attention_relative_position_bias,
                         )
                         if isinstance(hidden_states, tuple):
                             pass