Add original approach for cases with < INT_MAX elements in get_all_triplets_indices

Author: mkmenta

src/pytorch_metric_learning/utils/loss_and_miner_utils.py

@@ -86,6 +86,13 @@ def neg_pairs_from_tuple(indices_tuple):
 
 def get_all_triplets_indices(labels, ref_labels=None):
     all_matches, all_diffs = get_matches_and_diffs(labels, ref_labels)
+
+    if (all_matches.shape[0] * all_matches.shape[1] * all_matches.shape[1]
+            < torch.iinfo(torch.int32).max):
+        # torch.nonzero is not supported for tensors with more than INT_MAX elements
+        triplets = all_matches.unsqueeze(2) * all_diffs.unsqueeze(1)
+        return torch.where(triplets)
+
     all_matches, all_diffs = all_matches.bool(), all_diffs.bool()
 
     # Find anchors with at least a positive and a negative