Enable AutoModel loading

embeddings.py

@@ -105,8 +105,8 @@ class RotaryPositionalEncoding(nn.Module):
         seq_len = x.shape[1]
         cos: torch.Tensor = self.cos_cached  # type: ignore[assignment]
         sin: torch.Tensor = self.sin_cached  # type: ignore[assignment]
-        if seq_len > cos.shape[0]:
-            self._build_cache(seq_len)
+        if seq_len > cos.shape[0] or not torch.isfinite(cos[:seq_len]).all():
+            self._build_cache(max(seq_len, cos.shape[0]))
             cos = self.cos_cached  # type: ignore[assignment]
             sin = self.sin_cached  # type: ignore[assignment]
         return cos[:seq_len], sin[:seq_len]