zero shot experiment (fix v4)

model.py

@@ -262,7 +262,7 @@ class SentenceCoseNet(PreTrainedModel):
                 Similarity scores of shape (B, S, S)
         """
         # Concatenate embeddings (B, S, 2, D)
-        embeddings = torch.stack([embeddings_1, embeddings_2], dim=2)
+        embeddings = torch.stack([embeddings_1, embeddings_2], dim=-2)
         # Compute distances (B, S, 2, 2):
         embeddings = self.model.distance_layer(embeddings)
         # Return cosine similarities (B, S):