fix docstrings

model.py

@@ -232,14 +232,14 @@ class SentenceCoseNet(PreTrainedModel):
 
         Returns:
             torch.Tensor:
-                Similarity scores of shape (B, S, S)
+                Similarity scores of shape (B, S)
         """
         # Concatenate embeddings (B, S, 2, D)
         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):
-        return embeddings[..., 0, 1]
+        return (embeddings[..., 0, 1] + embeddings[..., 1, 0]) / 2
 
     def forward(
         self,