zero shot experiment (fix v3)

model.py

@@ -169,8 +169,8 @@ class SentenceCoseNet(PreTrainedModel):
                 `(batch_size, sequence_length, emb_dim)`.
         """
         # Convert to type:
-        x = input_ids.int()
-        mask = attention_mask if attention_mask is not None else None
+        x = input_ids.int().unsqueeze(1)
+        mask = attention_mask.unsqueeze(1) if attention_mask is not None else None
 
         # Embedding and positional encoding:
         x = self.model.embedding(x)
@@ -188,7 +188,7 @@ class SentenceCoseNet(PreTrainedModel):
 
         # Reshape x and mask:
         x = x.reshape(_b, _s, _t, _d)
-        return x.squeeze(0)
+        return x.squeeze(1)
 
     def get_sentence_embedding(
             self,
@@ -213,8 +213,8 @@ class SentenceCoseNet(PreTrainedModel):
                 Sentence embeddings of shape (B, D)
         """
         # Convert to type:
-        x = input_ids.int()
-        mask = attention_mask if attention_mask is not None else None
+        x = input_ids.int().unsqueeze(1)
+        mask = attention_mask.unsqueeze(1) if attention_mask is not None else None
 
         # Embedding and positional encoding:
         x = self.model.embedding(x)
@@ -242,7 +242,7 @@ class SentenceCoseNet(PreTrainedModel):
         # Apply normalization if required:
         if normalize:
             x = torch.nn.functional.normalize(x, p=2, dim=-1)
-        return x.squeeze(0)
+        return x.squeeze(1)
 
     def similarity(self, embeddings_1: torch.Tensor, embeddings_2: torch.Tensor) -> torch.Tensor:
         """