fix huggingface model weight missmatch

model.py

@@ -4,9 +4,6 @@
 # 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
@@ -291,9 +288,9 @@ class SentenceCoseNet(PreTrainedModel):
         """
         # Convert to type:
         if len(input_ids.shape) == 2:
-            x = input_ids.int().unsqueeze(1)
-            mask = attention_mask.unsqueeze(1) if attention_mask is not None else None
-            output = self.model(x=x, mask=mask).squeeze(1)
+            x = input_ids.int().unsqueeze(0)
+            mask = attention_mask.unsqueeze(0) if attention_mask is not None else None
+            output = self.model(x=x, mask=mask).squeeze(0)
         elif len(input_ids.shape) == 3:
             x = input_ids.int()
             mask = attention_mask if attention_mask is not None else None