Fixed check_shape test

Author: KevinMusgrave

src/pytorch_metric_learning/distances/base_distance.py

@@ -16,7 +16,7 @@ def __init__(
 
     def forward(self, query_emb, ref_emb=None):
         self.reset_stats()
-        self.check_embeddings_ndim(query_emb, ref_emb)
+        self.check_shapes(query_emb, ref_emb)
         query_emb_normalized = self.maybe_normalize(query_emb)
         if ref_emb is None:
             ref_emb = query_emb
@@ -89,7 +89,7 @@ def set_stats(self, stats_dict):
             self.add_to_recordable_attributes(name=k, is_stat=True)
             setattr(self, k, v)
 
-    def check_embeddings_ndim(self, query_emb, ref_emb):
+    def check_shapes(self, query_emb, ref_emb):
         if query_emb.ndim != 2 or (ref_emb is not None and ref_emb.ndim != 2):
             raise ValueError(
                 "embeddings must be a 2D tensor of shape (batch_size, embedding_size)"

tests/distances/test_custom_check_shape.py

@@ -0,0 +1,27 @@
+import unittest
+
+import torch
+
+from pytorch_metric_learning.distances import BaseDistance, LpDistance
+from pytorch_metric_learning.losses import TripletMarginLoss
+
+
+class CustomDistance(BaseDistance):
+    def compute_mat(self, query_emb, ref_emb):
+        return torch.randn(query_emb.shape[0], ref_emb.shape[0])
+
+    def check_shapes(self, query_emb, ref_emb):
+        pass
+
+
+class TestCustomCheckShape(unittest.TestCase):
+    def test_custom_embedding_ndim(self):
+        embeddings = torch.randn(32, 3, 128)
+        labels = torch.randint(0, 10, size=(32,))
+        loss_fn = TripletMarginLoss(distance=LpDistance())
+
+        with self.assertRaises(ValueError):
+            loss_fn(embeddings, labels)
+
+        loss_fn = TripletMarginLoss(distance=CustomDistance())
+        loss_fn(embeddings, labels)

tests/distances/test_custom_embedding_ndim.py

@@ -69,7 +69,7 @@ def test_collect_stats_flag(self):
 
     def test_check_shapes(self):
         embeddings = torch.randn(32, 512, 3)
-        labels = torch.randn(32)
+        labels = torch.randint(0, 10, size=(32,))
         loss_fn = TripletMarginLoss()
 
         # embeddings is 3-dimensional