Uploading patch

modeling_gpt_bert.py

@@ -138,7 +138,7 @@ class Attention(nn.Module):
             - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
         position_indices: torch.Tensor = self.make_log_bucket_position(position_indices, config.position_bucket_size, config.max_position_embeddings)
         position_indices = config.position_bucket_size - 1 + position_indices
-        self.register_buffer("position_indices", position_indices, persistent=True)
+        self.register_buffer("position_indices", position_indices, persistent=False)
 
         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
         self.scale: float = 1.0 / math.sqrt(3 * self.head_size)
@@ -301,18 +301,17 @@ class GPTBERT(GPTBERTPreTrainedModel):
         batch_size, seq_length = input_shape
 
         if attention_mask is None:
-            attention_mask = input_ids.new_ones((seq_length, seq_length), dtype=torch.bool).triu(diagonal=1).unsqueeze(0).unsqueeze(0)
-
-        if attention_mask is not None:
+            attention_mask = input_ids.new_zeros((batch_size, seq_length), dtype=torch.bool).unsqueeze(1).unsqueeze(2)
+        else:
             attention_mask = ~attention_mask.bool()
 
-            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)
+        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)
 
-            if self.is_causal:
-                attention_mask = attention_mask | input_ids.new_ones((seq_length, seq_length), dtype=torch.bool).triu(1).unsqueeze(0).unsqueeze(0)
+        if self.is_causal:
+            attention_mask = attention_mask | input_ids.new_ones((seq_length, seq_length), dtype=torch.bool).triu(1).unsqueeze(0).unsqueeze(0)
 
         static_embeddings, relative_embeddings = self.embedding(input_ids.t())
         contextualized_embeddings = [static_embeddings]