Update modeling_bert.py

modeling_bert.py

@@ -699,6 +699,11 @@ 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
@@ -766,7 +771,24 @@ class JinaBertEncoder(nn.Module):
 
         # Add alibi matrix to extended_attention_mask
         _, seqlen, _ = hidden_states.size()
-        alibi_bias = self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device).to(hidden_states.dtype)
+        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]
+        
         if self.gradient_checkpointing and self.training:
             if use_cache:
                 logger.warning_once(