fix positional encoding error

model.py

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

src/model/transformers/positional_encoding.py

@@ -56,7 +56,7 @@ class PositionalEncoding(torch.nn.Module):
         torch.Tensor
             Tensor of the same shape as the input with positional encodings added.
         """
-        return x + self.positional_encoding[:, :x.size(1), :]
+        return x + self.positional_encoding[:, :x.size(-2), :]
 # - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #
 #                        END OF FILE                        #
 # - x - x - x - x - x - x - x - x - x - x - x - x - x - x - #