Fix shape check

adapter-hub · Jan 26, 2022 · b83a66c · b83a66c
1 parent bf04239
commit b83a66c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/transformers/models/bart/modeling_bart.py b/src/transformers/models/bart/modeling_bart.py
@@ -1320,7 +1320,7 @@ def forward(
         )
         # sequence classification based on last token in sequence
         x = outputs[0]  # last hidden state
-        if input_ids is not None and x.shape == input_ids.shape:
+        if input_ids is not None and x.shape[1] == input_ids.shape[1]:
             eos_mask = input_ids.eq(self.config.eos_token_id)
             (eos_mask,) = adjust_tensors_for_parallel(x, eos_mask)
             if len(torch.unique(eos_mask.sum(1))) > 1:

diff --git a/src/transformers/models/mbart/modeling_mbart.py b/src/transformers/models/mbart/modeling_mbart.py
@@ -1320,7 +1320,7 @@ def forward(
         )
         # sequence classification based on last token in sequence
         x = outputs[0]  # last hidden state
-        if input_ids is not None and x.shape == input_ids.shape:
+        if input_ids is not None and x.shape[1] == input_ids.shape[1]:
             eos_mask = input_ids.eq(self.config.eos_token_id)
             (eos_mask,) = adjust_tensors_for_parallel(x, eos_mask)
             if len(torch.unique(eos_mask.sum(1))) > 1: