Update modeling_ltgbert.py

modeling_ltgbert.py

@@ -203,7 +203,7 @@ class Attention(nn.Module):
             - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
         position_indices = 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 = 1.0 / math.sqrt(3 * self.head_size)
@@ -234,7 +234,7 @@ class Attention(nn.Module):
             position_indices = torch.arange(query_len, dtype=torch.long).unsqueeze(1) \
                 - torch.arange(query_len, dtype=torch.long).unsqueeze(0)
             position_indices = self.make_log_bucket_position(position_indices, self.config.position_bucket_size, 512)
-            position_indices = self.position_bucket_size - 1 + position_indices
+            position_indices = self.config.position_bucket_size - 1 + position_indices
             self.position_indices = position_indices.to(hidden_states.device)
 
         hidden_states = self.pre_layer_norm(hidden_states)