Update modeling_nort5.py

modeling_nort5.py

@@ -405,7 +405,7 @@ class NorT5Model(NorT5PreTrainedModel):
     def get_decoder_output(
         self, target_ids, encoder_output, attention_mask
     ):
-        batch_size, seq_length = target_ids.shape
+        batch_size, seq_length, _ = encoder_output.shape
         device = target_ids.device
 
         if attention_mask is None: