Fix NaNs in model output

model.py

@@ -207,7 +207,16 @@ class NeoBERTPreTrainedModel(PreTrainedModel):
             module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
         elif isinstance(module, nn.Embedding):
             module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)
+        elif isinstance(module, NeoBERT):
+            freqs = precompute_freqs_cis(
+                self.config.hidden_size // self.config.num_attention_heads,
+                self.config.max_length
+            )
 
+            if module.freqs_cis.device.type == "meta":
+                module.freqs_cis = freqs.to(module.freqs_cis.device)
+            else:
+                module.freqs_cis.copy_(freqs.to(device=module.freqs_cis.device, dtype=module.freqs_cis.dtype))
 
 class NeoBERT(NeoBERTPreTrainedModel):
     config_class = NeoBERTConfig