Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -619,7 +619,9 @@ class NomicBertRotaryEmbedding(nn.Module):
         Apply rotary embedding *inplace* to qkv and / or kv.
         """
         seqlen = qkv.shape[1]
-        if max_seqlen is not None:
+        if seqlen > self._seq_len_cached:
+            self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+        elif max_seqlen is not None:
             self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
         elif isinstance(seqlen_offset, int):
             self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype)