Fixes bug where BaseDistance.initial_avg_query_norm was not actually being set

Author: KevinMusgrave

src/pytorch_metric_learning/__init__.py

@@ -1 +1 @@
-__version__ = "2.1.0"
+__version__ = "2.1.1"

src/pytorch_metric_learning/distances/base_distance.py

@@ -79,7 +79,9 @@ def set_default_stats(
     ):
         if self.collect_stats:
             with torch.no_grad():
-                self.initial_avg_query_norm: torch.mean(self.get_norm(query_emb)).item()
+                self.initial_avg_query_norm = torch.mean(
+                    self.get_norm(query_emb)
+                ).item()
                 self.initial_avg_ref_norm = torch.mean(self.get_norm(ref_emb)).item()
                 self.final_avg_query_norm = torch.mean(
                     self.get_norm(query_emb_normalized)

tests/distances/test_collected_stats.py

@@ -0,0 +1,20 @@
+import unittest
+
+import torch
+
+from pytorch_metric_learning.distances import LpDistance
+
+from .. import WITH_COLLECT_STATS
+
+
+class TestCollectedStats(unittest.TestCase):
+    @unittest.skipUnless(WITH_COLLECT_STATS, "WITH_COLLECT_STATS is false")
+    def test_collected_stats(self):
+        x = torch.randn(32, 128)
+        d = LpDistance()
+        d(x)
+
+        self.assertNotEqual(d.initial_avg_query_norm, 0)
+        self.assertNotEqual(d.initial_avg_ref_norm, 0)
+        self.assertNotEqual(d.final_avg_query_norm, 0)
+        self.assertNotEqual(d.final_avg_ref_norm, 0)