Minor memory handling fixes

modeling_ltgbert.py

@@ -230,7 +230,7 @@ class Attention(nn.Module):
 
         query = query.reshape(query_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
         key = key.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
-        value = value.view(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
+        value = value.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
 
         attention_scores = torch.bmm(query, key.transpose(1, 2) * self.scale)
 
@@ -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)