Moved ndim check into BaseDistance

Author: KevinMusgrave

src/pytorch_metric_learning/__init__.py

@@ -1 +1 @@
-__version__ = "1.7.1"
+__version__ = "1.7.2"

src/pytorch_metric_learning/distances/base_distance.py

@@ -16,6 +16,7 @@ def __init__(
 
     def forward(self, query_emb, ref_emb=None):
         self.reset_stats()
+        self.check_embeddings_ndim(query_emb, ref_emb)
         query_emb_normalized = self.maybe_normalize(query_emb)
         if ref_emb is None:
             ref_emb = query_emb
@@ -87,3 +88,9 @@ def set_stats(self, stats_dict):
         for k, v in stats_dict.items():
             self.add_to_recordable_attributes(name=k, is_stat=True)
             setattr(self, k, v)
+
+    def check_embeddings_ndim(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)"
+            )

src/pytorch_metric_learning/utils/common_functions.py

@@ -412,10 +412,6 @@ def return_input(x):
 def check_shapes(embeddings, labels):
     if labels is not None and embeddings.shape[0] != labels.shape[0]:
         raise ValueError("Number of embeddings must equal number of labels")
-    if embeddings.ndim != 2:
-        raise ValueError(
-            "embeddings must be a 2D tensor of shape (batch_size, embedding_size)"
-        )
     if labels is not None and labels.ndim != 1:
         raise ValueError("labels must be a 1D tensor of shape (batch_size,)")
 

tests/distances/test_custom_embedding_ndim.py

@@ -0,0 +1,23 @@
+import unittest
+
+import torch
+
+from pytorch_metric_learning.distances import BaseDistance, LpDistance
+
+
+class CustomDistance(BaseDistance):
+    def compute_mat(self, query_emb, ref_emb):
+        pass
+
+    def check_embeddings_ndim(self, query_emb, ref_emb):
+        pass
+
+
+class TestCustomEmbeddingNdim(unittest.TestCase):
+    def test_custom_embedding_ndim(self):
+        embeddings = torch.randn(32, 3, 128)
+
+        dist_fn = LpDistance()
+
+        with self.assertRaises(ValueError):
+            dist_fn(embeddings)