fix huggingface model weight missmatch

model.py

@@ -4,6 +4,9 @@
 # Universidad de Alcalá - Escuela Politécnica Superior      #
 #                                                           #
 # - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
+import os
+from safetensors.torch import load_file
+
 import torch
 from transformers import PreTrainedModel, PretrainedConfig
 from src.model import SegmentationNetwork
@@ -138,7 +141,13 @@ class SentenceCoseNet(PreTrainedModel):
         
         # Core PyTorch model
         self.model = SegmentationNetwork(self.to_model_config(config))
-        state_dict = torch.load("model.safetensors", map_location="cpu")
+
+        weights_path = os.path.join(
+            config._name_or_path,
+            "model.safetensors"
+        )
+
+        state_dict = load_file(weights_path)
         self.model.load_state_dict(state_dict, strict=True)
         self.model.eval()