formatted code

KevinMusgrave · Jul 24, 2024 · ef1bd06 · ef1bd06
1 parent 0db2483
commit ef1bd06
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/pytorch_metric_learning/utils/distributed.py b/src/pytorch_metric_learning/utils/distributed.py
@@ -109,7 +109,14 @@ def forward(
             return self.loss(embeddings, labels, indices_tuple, ref_emb, ref_labels)
 
         world_size = torch.distributed.get_world_size()
-        common_args = [embeddings, labels, indices_tuple, ref_emb, ref_labels, world_size]
+        common_args = [
+            embeddings,
+            labels,
+            indices_tuple,
+            ref_emb,
+            ref_labels,
+            world_size,
+        ]
         if isinstance(self.loss, CrossBatchMemory):
             return self.forward_cross_batch(*common_args, enqueue_mask)
         return self.forward_regular_loss(*common_args)