Gradient in distributed training
Here is the current implementation of the distributed wrapper.

https://github.com/KevinMusgrave/pytorch-metric-learning/blob/6bfa880b8d2acafb7c6d52041d2bb14ed41aee76/src/pytorch_metric_learning/utils/distributed.py#L35-L59

But according to this [blog](https://amsword.medium.com/gradient-backpropagation-with-torch-distributed-all-gather-9f3941a381f8), the loss function should be multiplied by `world_size`. 

What do you think about it? Maybe I can create a PR to fix it.
diff --git a/src/pytorch_metric_learning/__init__.py b/src/pytorch_metric_learning/__init__.py
index a7be658..f2bdfbe 100644
--- a/src/pytorch_metric_learning/__init__.py
+++ b/src/pytorch_metric_learning/__init__.py
@@ -1 +1 @@
-__version__ = "1.0.0.dev2"
+__version__ = "1.0.0.dev3"
diff --git a/src/pytorch_metric_learning/utils/distributed.py b/src/pytorch_metric_learning/utils/distributed.py
index e6108c3..9e11955 100644
--- a/src/pytorch_metric_learning/utils/distributed.py
+++ b/src/pytorch_metric_learning/utils/distributed.py
@@ -56,7 +56,8 @@ class DistributedLossWrapper(torch.nn.Module):
 
     def forward(self, embeddings, labels, *args, **kwargs):
         embeddings, labels = all_gather_embeddings_labels(embeddings, labels)
-        return self.loss(embeddings, labels, *args, **kwargs)
+        world_size = torch.distributed.get_world_size()
+        return world_size * self.loss(embeddings, labels, *args, **kwargs)
 
 
 class DistributedMinerWrapper(torch.nn.Module):
diff --git a/tests/utils/test_distributed.py b/tests/utils/test_distributed.py
index 8905bcd..9878704 100644
--- a/tests/utils/test_distributed.py
+++ b/tests/utils/test_distributed.py
@@ -132,7 +132,7 @@ def single_process_function(
             correct_loss_size = original_loss_fn.reducer.losses_size
             assert loss_size == correct_loss_size
 
-    assert torch.isclose(loss, torch.from_numpy(correct_loss).to(device))
+    assert torch.isclose(loss, torch.from_numpy(correct_loss).to(device) * world_size)
     assert miner_fn.miner.num_pos_pairs == original_miner_fn.num_pos_pairs
     assert miner_fn.miner.num_neg_pairs == original_miner_fn.num_neg_pairs
     for i in range(len(correct_indices_tuple)):
@@ -245,10 +245,7 @@ class TestDistributedLossWrapper(unittest.TestCase):
                     all_outputs, all_labels, correct_indices_tuple
                 )
 
-                if TEST_DEVICE == torch.device("cpu"):
-                    correct_loss.backward(retain_graph=True)
-                else:
-                    (correct_loss / world_size).backward(retain_graph=True)
+                correct_loss.backward(retain_graph=True)
                 optimizer.step()
                 if not is_tuple_loss:
                     for p in original_loss_fn.parameters():
['tests/utils/test_distributed.py', 'src/pytorch_metric_learning/utils/distributed.py', 'src/pytorch_metric_learning/__init__.py']