Update modeling_nort5.py

modeling_nort5.py

@@ -221,7 +221,7 @@ class Attention(nn.Module):
             - torch.arange(512, dtype=torch.long).unsqueeze(0)
         position_indices = self.make_log_bucket_position(position_indices, config.position_bucket_size, 512)
         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 = 1.0 / math.sqrt(3 * self.head_size)
@@ -271,7 +271,7 @@ class Attention(nn.Module):
                 - torch.arange(max(query_len, key_len), dtype=torch.long).unsqueeze(0)
             position_indices = self.make_log_bucket_position(position_indices, self.config.position_bucket_size, 512)
             position_indices = self.config.position_bucket_size - 1 + position_indices
-            self.register_buffer("position_indices", position_indices.to(q.device), persistent=True)
+            self.register_buffer("position_indices", position_indices.to(q.device), persistent=False)
 
         q = self.pre_layer_norm(q)
         query = self.in_proj_q(q)  # shape: [T, B, D]