Update modeling_ltgbert.py

modeling_ltgbert.py

@@ -346,7 +346,10 @@ class LtgbertModel(LtgbertPreTrainedModel):
         if self.config.is_decoder:
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) | torch.triu(torch.ones(seq_length, seq_length, dtype=torch.bool, device=device), 1).unsqueeze(0).unsqueeze(0)
         else:
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            if len(attention_mask.size()) == 2:
+                attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            elif len(attention_mask.size()) == 3:
+                attention_mask = attention_mask.unsqueeze(1)
  
         static_embeddings, relative_embedding = self.embedding(input_ids.t())
         contextualized_embeddings, attention_probs = self.transformer(static_embeddings, attention_mask, relative_embedding)