Refactor functions/losses.py

prototorch/functions/losses.py

@@ -3,12 +3,19 @@
 import torch
 
 
-def _get_dp_dm(distances, targets, plabels):
-    matcher = torch.eq(targets.unsqueeze(dim=1), plabels)
-    if plabels.ndim == 2:
+def _get_matcher(targets, labels):
+    """Returns a boolean tensor."""
+    matcher = torch.eq(targets.unsqueeze(dim=1), labels)
+    if labels.ndim == 2:
         # if the labels are one-hot vectors
         nclasses = targets.size()[1]
         matcher = torch.eq(torch.sum(matcher, dim=-1), nclasses)
+    return matcher
+
+
+def _get_dp_dm(distances, targets, plabels):
+    """Returns the d+ and d- values for a batch of distances."""
+    matcher = _get_matcher(targets, plabels)
     not_matcher = torch.bitwise_not(matcher)
 
     inf = torch.full_like(distances, fill_value=float("inf"))