Update modeling_bert.py

modeling_bert.py

@@ -697,11 +697,6 @@ class JinaBertEncoder(nn.Module):
         )
         self.gradient_checkpointing = False
         self.num_attention_heads = config.num_attention_heads
-        self.register_buffer(
-            "alibi",
-            self.rebuild_alibi_tensor(size=config.max_position_embeddings),
-            persistent=False,
-        )
 
     def rebuild_alibi_tensor(
         self, size: int, device: Optional[Union[torch.device, str]] = None
@@ -769,23 +764,7 @@ class JinaBertEncoder(nn.Module):
 
         # Add alibi matrix to extended_attention_mask
         _, seqlen, _ = hidden_states.size()
-        if self._current_alibi_size < seqlen:
-            # Rebuild the alibi tensor when needed
-            warnings.warn(
-                f'Increasing alibi size from {self._current_alibi_size} to {seqlen}.'
-            )
-            self.register_buffer(
-                "alibi",
-                self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device).to(
-                    hidden_states.dtype
-                ),
-                persistent=False,
-            )
-        elif self.alibi.device != hidden_states.device:
-            # Device catch-up
-            self.alibi = self.alibi.to(hidden_states.device)
-
-        alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
+        alibi_bias = self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device).to(hidden_states.dtype)
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(