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)
@@ -302,14 +302,13 @@ class GPTBERT(GPTBERTPreTrainedModel):
 
         if attention_mask is None:
             attention_mask = input_ids.new_zeros((batch_size, seq_length), dtype=torch.bool).unsqueeze(1).unsqueeze(2)
-
-        if attention_mask is not None:
+        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)