fix api (#11337)

Signed-off-by: yaoyu-33 <[email protected]>
NVIDIA · Nov 20, 2024 · 4e56c40 · 4e56c40
1 parent 8b7999d
commit 4e56c40
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/nemo/collections/vision/modules/common/megatron/vision_transformer.py b/nemo/collections/vision/modules/common/megatron/vision_transformer.py
@@ -169,6 +169,8 @@ def forward(
         self_attention_relative_position_bias=None,
         cross_attention_relative_position_bias=None,
         checkpoint_core_attention=False,
+        decoder_max_sequence_len=None,
+        encoder_max_sequence_len=None,
     ):
         # Self attention.
         if rotary_pos_emb is not None:
@@ -373,6 +375,8 @@ def forward(
         self_attention_relative_position_bias=None,
         cross_attention_relative_position_bias=None,
         checkpoint_core_attention=False,
+        decoder_max_sequence_len=None,
+        encoder_max_sequence_len=None,
     ):
         kwargs = locals()
         for key in ["self", "__class__"]: