We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
CUDA errors in NTXentLoss with gloo backend in multi-gpu training 
I was wondering if the `gloo` distributed communication package for multi-gpu training is officially supported by lightly. It seems like e.g. NTXentLoss doesn't work with `gloo` (I'm using pytorch lightning): I get CUDA errors, even when setting `gather_distributed = False`.
I can fix the issue when using `gather_distributed = False` by replacing the line
https://github.com/lightly-ai/lightly/blob/master/lightly/loss/ntx_ent_loss.py#L164
by
```python
labels = labels + batch_size * (dist.rank() if gather_distributed else 0)
```
but then of course I can't use `gather_distributed = True` anymore. 

Using the `nccl` backend, everything works fine, but `nccl` is not working well on some of our machines, so unfortunately I'm stuck with `gloo`. I think using `gloo` might be too much of an exception to fix the problem for `gather_distributed = True`, but maybe it'd help to just replace the line above and mention somewhere in the documentation that `gather_distributed` is not supported for `gloo`?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `lightly/loss/ntx_ent_loss.py`
Content:
```
1 """ Contrastive Loss Functions """
2 
3 # Copyright (c) 2020. Lightly AG and its affiliates.
4 # All Rights Reserved
5 
6 import torch
7 from torch import nn
8 
9 from lightly.loss.memory_bank import MemoryBankModule
10 from lightly.utils import dist
11 
12 
13 class NTXentLoss(MemoryBankModule):
14     """Implementation of the Contrastive Cross Entropy Loss.
15 
16     This implementation follows the SimCLR[0] paper. If you enable the memory
17     bank by setting the `memory_bank_size` value > 0 the loss behaves like
18     the one described in the MoCo[1] paper.
19 
20     - [0] SimCLR, 2020, https://arxiv.org/abs/2002.05709
21     - [1] MoCo, 2020, https://arxiv.org/abs/1911.05722
22 
23     Attributes:
24         temperature:
25             Scale logits by the inverse of the temperature.
26         memory_bank_size:
27             Number of negative samples to store in the memory bank.
28             Use 0 for SimCLR. For MoCo we typically use numbers like 4096 or 65536.
29         gather_distributed:
30             If True then negatives from all gpus are gathered before the
31             loss calculation. This flag has no effect if memory_bank_size > 0.
32 
33     Raises:
34         ValueError: If abs(temperature) < 1e-8 to prevent divide by zero.
35 
36     Examples:
37 
38         >>> # initialize loss function without memory bank
39         >>> loss_fn = NTXentLoss(memory_bank_size=0)
40         >>>
41         >>> # generate two random transforms of images
42         >>> t0 = transforms(images)
43         >>> t1 = transforms(images)
44         >>>
45         >>> # feed through SimCLR or MoCo model
46         >>> batch = torch.cat((t0, t1), dim=0)
47         >>> output = model(batch)
48         >>>
49         >>> # calculate loss
50         >>> loss = loss_fn(output)
51 
52     """
53 
54     def __init__(
55         self,
56         temperature: float = 0.5,
57         memory_bank_size: int = 0,
58         gather_distributed: bool = False,
59     ):
60         super(NTXentLoss, self).__init__(size=memory_bank_size)
61         self.temperature = temperature
62         self.gather_distributed = gather_distributed
63         self.cross_entropy = nn.CrossEntropyLoss(reduction="mean")
64         self.eps = 1e-8
65 
66         if abs(self.temperature) < self.eps:
67             raise ValueError(
68                 "Illegal temperature: abs({}) < 1e-8".format(self.temperature)
69             )
70 
71     def forward(self, out0: torch.Tensor, out1: torch.Tensor):
72         """Forward pass through Contrastive Cross-Entropy Loss.
73 
74         If used with a memory bank, the samples from the memory bank are used
75         as negative examples. Otherwise, within-batch samples are used as
76         negative samples.
77 
78         Args:
79             out0:
80                 Output projections of the first set of transformed images.
81                 Shape: (batch_size, embedding_size)
82             out1:
83                 Output projections of the second set of transformed images.
84                 Shape: (batch_size, embedding_size)
85 
86         Returns:
87             Contrastive Cross Entropy Loss value.
88 
89         """
90 
91         device = out0.device
92         batch_size, _ = out0.shape
93 
94         # normalize the output to length 1
95         out0 = nn.functional.normalize(out0, dim=1)
96         out1 = nn.functional.normalize(out1, dim=1)
97 
98         # ask memory bank for negative samples and extend it with out1 if
99         # out1 requires a gradient, otherwise keep the same vectors in the
100         # memory bank (this allows for keeping the memory bank constant e.g.
101         # for evaluating the loss on the test set)
102         # out1: shape: (batch_size, embedding_size)
103         # negatives: shape: (embedding_size, memory_bank_size)
104         out1, negatives = super(NTXentLoss, self).forward(
105             out1, update=out0.requires_grad
106         )
107 
108         # We use the cosine similarity, which is a dot product (einsum) here,
109         # as all vectors are already normalized to unit length.
110         # Notation in einsum: n = batch_size, c = embedding_size and k = memory_bank_size.
111 
112         if negatives is not None:
113             # use negatives from memory bank
114             negatives = negatives.to(device)
115 
116             # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity
117             # of the i-th sample in the batch to its positive pair
118             sim_pos = torch.einsum("nc,nc->n", out0, out1).unsqueeze(-1)
119 
120             # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity
121             # of the i-th sample to the j-th negative sample
122             sim_neg = torch.einsum("nc,ck->nk", out0, negatives)
123 
124             # set the labels to the first "class", i.e. sim_pos,
125             # so that it is maximized in relation to sim_neg
126             logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature
127             labels = torch.zeros(logits.shape[0], device=device, dtype=torch.long)
128 
129         else:
130             # user other samples from batch as negatives
131             # and create diagonal mask that only selects similarities between
132             # views of the same image
133             if self.gather_distributed and dist.world_size() > 1:
134                 # gather hidden representations from other processes
135                 out0_large = torch.cat(dist.gather(out0), 0)
136                 out1_large = torch.cat(dist.gather(out1), 0)
137                 diag_mask = dist.eye_rank(batch_size, device=out0.device)
138             else:
139                 # single process
140                 out0_large = out0
141                 out1_large = out1
142                 diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool)
143 
144             # calculate similiarities
145             # here n = batch_size and m = batch_size * world_size
146             # the resulting vectors have shape (n, m)
147             logits_00 = torch.einsum("nc,mc->nm", out0, out0_large) / self.temperature
148             logits_01 = torch.einsum("nc,mc->nm", out0, out1_large) / self.temperature
149             logits_10 = torch.einsum("nc,mc->nm", out1, out0_large) / self.temperature
150             logits_11 = torch.einsum("nc,mc->nm", out1, out1_large) / self.temperature
151 
152             # remove simliarities between same views of the same image
153             logits_00 = logits_00[~diag_mask].view(batch_size, -1)
154             logits_11 = logits_11[~diag_mask].view(batch_size, -1)
155 
156             # concatenate logits
157             # the logits tensor in the end has shape (2*n, 2*m-1)
158             logits_0100 = torch.cat([logits_01, logits_00], dim=1)
159             logits_1011 = torch.cat([logits_10, logits_11], dim=1)
160             logits = torch.cat([logits_0100, logits_1011], dim=0)
161 
162             # create labels
163             labels = torch.arange(batch_size, device=device, dtype=torch.long)
164             labels = labels + dist.rank() * batch_size
165             labels = labels.repeat(2)
166 
167         loss = self.cross_entropy(logits, labels)
168 
169         return loss
170
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/lightly/loss/ntx_ent_loss.py b/lightly/loss/ntx_ent_loss.py
--- a/lightly/loss/ntx_ent_loss.py
+++ b/lightly/loss/ntx_ent_loss.py
@@ -161,7 +161,8 @@
 
             # create labels
             labels = torch.arange(batch_size, device=device, dtype=torch.long)
-            labels = labels + dist.rank() * batch_size
+            if self.gather_distributed:
+                labels = labels + dist.rank() * batch_size
             labels = labels.repeat(2)
 
         loss = self.cross_entropy(logits, labels)

{"golden_diff": "diff --git a/lightly/loss/ntx_ent_loss.py b/lightly/loss/ntx_ent_loss.py\n--- a/lightly/loss/ntx_ent_loss.py\n+++ b/lightly/loss/ntx_ent_loss.py\n@@ -161,7 +161,8 @@\n \n             # create labels\n             labels = torch.arange(batch_size, device=device, dtype=torch.long)\n-            labels = labels + dist.rank() * batch_size\n+            if self.gather_distributed:\n+                labels = labels + dist.rank() * batch_size\n             labels = labels.repeat(2)\n \n         loss = self.cross_entropy(logits, labels)\n", "issue": "CUDA errors in NTXentLoss with gloo backend in multi-gpu training \nI was wondering if the `gloo` distributed communication package for multi-gpu training is officially supported by lightly. It seems like e.g. NTXentLoss doesn't work with `gloo` (I'm using pytorch lightning): I get CUDA errors, even when setting `gather_distributed = False`.\r\nI can fix the issue when using `gather_distributed = False` by replacing the line\r\nhttps://github.com/lightly-ai/lightly/blob/master/lightly/loss/ntx_ent_loss.py#L164\r\nby\r\n```python\r\nlabels = labels + batch_size * (dist.rank() if gather_distributed else 0)\r\n```\r\nbut then of course I can't use `gather_distributed = True` anymore. \r\n\r\nUsing the `nccl` backend, everything works fine, but `nccl` is not working well on some of our machines, so unfortunately I'm stuck with `gloo`. I think using `gloo` might be too much of an exception to fix the problem for `gather_distributed = True`, but maybe it'd help to just replace the line above and mention somewhere in the documentation that `gather_distributed` is not supported for `gloo`?\n", "before_files": [{"content": "\"\"\" Contrastive Loss Functions \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nfrom torch import nn\n\nfrom lightly.loss.memory_bank import MemoryBankModule\nfrom lightly.utils import dist\n\n\nclass NTXentLoss(MemoryBankModule):\n    \"\"\"Implementation of the Contrastive Cross Entropy Loss.\n\n    This implementation follows the SimCLR[0] paper. If you enable the memory\n    bank by setting the `memory_bank_size` value > 0 the loss behaves like\n    the one described in the MoCo[1] paper.\n\n    - [0] SimCLR, 2020, https://arxiv.org/abs/2002.05709\n    - [1] MoCo, 2020, https://arxiv.org/abs/1911.05722\n\n    Attributes:\n        temperature:\n            Scale logits by the inverse of the temperature.\n        memory_bank_size:\n            Number of negative samples to store in the memory bank.\n            Use 0 for SimCLR. For MoCo we typically use numbers like 4096 or 65536.\n        gather_distributed:\n            If True then negatives from all gpus are gathered before the\n            loss calculation. This flag has no effect if memory_bank_size > 0.\n\n    Raises:\n        ValueError: If abs(temperature) < 1e-8 to prevent divide by zero.\n\n    Examples:\n\n        >>> # initialize loss function without memory bank\n        >>> loss_fn = NTXentLoss(memory_bank_size=0)\n        >>>\n        >>> # generate two random transforms of images\n        >>> t0 = transforms(images)\n        >>> t1 = transforms(images)\n        >>>\n        >>> # feed through SimCLR or MoCo model\n        >>> batch = torch.cat((t0, t1), dim=0)\n        >>> output = model(batch)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(output)\n\n    \"\"\"\n\n    def __init__(\n        self,\n        temperature: float = 0.5,\n        memory_bank_size: int = 0,\n        gather_distributed: bool = False,\n    ):\n        super(NTXentLoss, self).__init__(size=memory_bank_size)\n        self.temperature = temperature\n        self.gather_distributed = gather_distributed\n        self.cross_entropy = nn.CrossEntropyLoss(reduction=\"mean\")\n        self.eps = 1e-8\n\n        if abs(self.temperature) < self.eps:\n            raise ValueError(\n                \"Illegal temperature: abs({}) < 1e-8\".format(self.temperature)\n            )\n\n    def forward(self, out0: torch.Tensor, out1: torch.Tensor):\n        \"\"\"Forward pass through Contrastive Cross-Entropy Loss.\n\n        If used with a memory bank, the samples from the memory bank are used\n        as negative examples. Otherwise, within-batch samples are used as\n        negative samples.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images.\n                Shape: (batch_size, embedding_size)\n            out1:\n                Output projections of the second set of transformed images.\n                Shape: (batch_size, embedding_size)\n\n        Returns:\n            Contrastive Cross Entropy Loss value.\n\n        \"\"\"\n\n        device = out0.device\n        batch_size, _ = out0.shape\n\n        # normalize the output to length 1\n        out0 = nn.functional.normalize(out0, dim=1)\n        out1 = nn.functional.normalize(out1, dim=1)\n\n        # ask memory bank for negative samples and extend it with out1 if\n        # out1 requires a gradient, otherwise keep the same vectors in the\n        # memory bank (this allows for keeping the memory bank constant e.g.\n        # for evaluating the loss on the test set)\n        # out1: shape: (batch_size, embedding_size)\n        # negatives: shape: (embedding_size, memory_bank_size)\n        out1, negatives = super(NTXentLoss, self).forward(\n            out1, update=out0.requires_grad\n        )\n\n        # We use the cosine similarity, which is a dot product (einsum) here,\n        # as all vectors are already normalized to unit length.\n        # Notation in einsum: n = batch_size, c = embedding_size and k = memory_bank_size.\n\n        if negatives is not None:\n            # use negatives from memory bank\n            negatives = negatives.to(device)\n\n            # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity\n            # of the i-th sample in the batch to its positive pair\n            sim_pos = torch.einsum(\"nc,nc->n\", out0, out1).unsqueeze(-1)\n\n            # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity\n            # of the i-th sample to the j-th negative sample\n            sim_neg = torch.einsum(\"nc,ck->nk\", out0, negatives)\n\n            # set the labels to the first \"class\", i.e. sim_pos,\n            # so that it is maximized in relation to sim_neg\n            logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature\n            labels = torch.zeros(logits.shape[0], device=device, dtype=torch.long)\n\n        else:\n            # user other samples from batch as negatives\n            # and create diagonal mask that only selects similarities between\n            # views of the same image\n            if self.gather_distributed and dist.world_size() > 1:\n                # gather hidden representations from other processes\n                out0_large = torch.cat(dist.gather(out0), 0)\n                out1_large = torch.cat(dist.gather(out1), 0)\n                diag_mask = dist.eye_rank(batch_size, device=out0.device)\n            else:\n                # single process\n                out0_large = out0\n                out1_large = out1\n                diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool)\n\n            # calculate similiarities\n            # here n = batch_size and m = batch_size * world_size\n            # the resulting vectors have shape (n, m)\n            logits_00 = torch.einsum(\"nc,mc->nm\", out0, out0_large) / self.temperature\n            logits_01 = torch.einsum(\"nc,mc->nm\", out0, out1_large) / self.temperature\n            logits_10 = torch.einsum(\"nc,mc->nm\", out1, out0_large) / self.temperature\n            logits_11 = torch.einsum(\"nc,mc->nm\", out1, out1_large) / self.temperature\n\n            # remove simliarities between same views of the same image\n            logits_00 = logits_00[~diag_mask].view(batch_size, -1)\n            logits_11 = logits_11[~diag_mask].view(batch_size, -1)\n\n            # concatenate logits\n            # the logits tensor in the end has shape (2*n, 2*m-1)\n            logits_0100 = torch.cat([logits_01, logits_00], dim=1)\n            logits_1011 = torch.cat([logits_10, logits_11], dim=1)\n            logits = torch.cat([logits_0100, logits_1011], dim=0)\n\n            # create labels\n            labels = torch.arange(batch_size, device=device, dtype=torch.long)\n            labels = labels + dist.rank() * batch_size\n            labels = labels.repeat(2)\n\n        loss = self.cross_entropy(logits, labels)\n\n        return loss\n", "path": "lightly/loss/ntx_ent_loss.py"}], "after_files": [{"content": "\"\"\" Contrastive Loss Functions \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nfrom torch import nn\n\nfrom lightly.loss.memory_bank import MemoryBankModule\nfrom lightly.utils import dist\n\n\nclass NTXentLoss(MemoryBankModule):\n    \"\"\"Implementation of the Contrastive Cross Entropy Loss.\n\n    This implementation follows the SimCLR[0] paper. If you enable the memory\n    bank by setting the `memory_bank_size` value > 0 the loss behaves like\n    the one described in the MoCo[1] paper.\n\n    - [0] SimCLR, 2020, https://arxiv.org/abs/2002.05709\n    - [1] MoCo, 2020, https://arxiv.org/abs/1911.05722\n\n    Attributes:\n        temperature:\n            Scale logits by the inverse of the temperature.\n        memory_bank_size:\n            Number of negative samples to store in the memory bank.\n            Use 0 for SimCLR. For MoCo we typically use numbers like 4096 or 65536.\n        gather_distributed:\n            If True then negatives from all gpus are gathered before the\n            loss calculation. This flag has no effect if memory_bank_size > 0.\n\n    Raises:\n        ValueError: If abs(temperature) < 1e-8 to prevent divide by zero.\n\n    Examples:\n\n        >>> # initialize loss function without memory bank\n        >>> loss_fn = NTXentLoss(memory_bank_size=0)\n        >>>\n        >>> # generate two random transforms of images\n        >>> t0 = transforms(images)\n        >>> t1 = transforms(images)\n        >>>\n        >>> # feed through SimCLR or MoCo model\n        >>> batch = torch.cat((t0, t1), dim=0)\n        >>> output = model(batch)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(output)\n\n    \"\"\"\n\n    def __init__(\n        self,\n        temperature: float = 0.5,\n        memory_bank_size: int = 0,\n        gather_distributed: bool = False,\n    ):\n        super(NTXentLoss, self).__init__(size=memory_bank_size)\n        self.temperature = temperature\n        self.gather_distributed = gather_distributed\n        self.cross_entropy = nn.CrossEntropyLoss(reduction=\"mean\")\n        self.eps = 1e-8\n\n        if abs(self.temperature) < self.eps:\n            raise ValueError(\n                \"Illegal temperature: abs({}) < 1e-8\".format(self.temperature)\n            )\n\n    def forward(self, out0: torch.Tensor, out1: torch.Tensor):\n        \"\"\"Forward pass through Contrastive Cross-Entropy Loss.\n\n        If used with a memory bank, the samples from the memory bank are used\n        as negative examples. Otherwise, within-batch samples are used as\n        negative samples.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images.\n                Shape: (batch_size, embedding_size)\n            out1:\n                Output projections of the second set of transformed images.\n                Shape: (batch_size, embedding_size)\n\n        Returns:\n            Contrastive Cross Entropy Loss value.\n\n        \"\"\"\n\n        device = out0.device\n        batch_size, _ = out0.shape\n\n        # normalize the output to length 1\n        out0 = nn.functional.normalize(out0, dim=1)\n        out1 = nn.functional.normalize(out1, dim=1)\n\n        # ask memory bank for negative samples and extend it with out1 if\n        # out1 requires a gradient, otherwise keep the same vectors in the\n        # memory bank (this allows for keeping the memory bank constant e.g.\n        # for evaluating the loss on the test set)\n        # out1: shape: (batch_size, embedding_size)\n        # negatives: shape: (embedding_size, memory_bank_size)\n        out1, negatives = super(NTXentLoss, self).forward(\n            out1, update=out0.requires_grad\n        )\n\n        # We use the cosine similarity, which is a dot product (einsum) here,\n        # as all vectors are already normalized to unit length.\n        # Notation in einsum: n = batch_size, c = embedding_size and k = memory_bank_size.\n\n        if negatives is not None:\n            # use negatives from memory bank\n            negatives = negatives.to(device)\n\n            # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity\n            # of the i-th sample in the batch to its positive pair\n            sim_pos = torch.einsum(\"nc,nc->n\", out0, out1).unsqueeze(-1)\n\n            # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity\n            # of the i-th sample to the j-th negative sample\n            sim_neg = torch.einsum(\"nc,ck->nk\", out0, negatives)\n\n            # set the labels to the first \"class\", i.e. sim_pos,\n            # so that it is maximized in relation to sim_neg\n            logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature\n            labels = torch.zeros(logits.shape[0], device=device, dtype=torch.long)\n\n        else:\n            # user other samples from batch as negatives\n            # and create diagonal mask that only selects similarities between\n            # views of the same image\n            if self.gather_distributed and dist.world_size() > 1:\n                # gather hidden representations from other processes\n                out0_large = torch.cat(dist.gather(out0), 0)\n                out1_large = torch.cat(dist.gather(out1), 0)\n                diag_mask = dist.eye_rank(batch_size, device=out0.device)\n            else:\n                # single process\n                out0_large = out0\n                out1_large = out1\n                diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool)\n\n            # calculate similiarities\n            # here n = batch_size and m = batch_size * world_size\n            # the resulting vectors have shape (n, m)\n            logits_00 = torch.einsum(\"nc,mc->nm\", out0, out0_large) / self.temperature\n            logits_01 = torch.einsum(\"nc,mc->nm\", out0, out1_large) / self.temperature\n            logits_10 = torch.einsum(\"nc,mc->nm\", out1, out0_large) / self.temperature\n            logits_11 = torch.einsum(\"nc,mc->nm\", out1, out1_large) / self.temperature\n\n            # remove simliarities between same views of the same image\n            logits_00 = logits_00[~diag_mask].view(batch_size, -1)\n            logits_11 = logits_11[~diag_mask].view(batch_size, -1)\n\n            # concatenate logits\n            # the logits tensor in the end has shape (2*n, 2*m-1)\n            logits_0100 = torch.cat([logits_01, logits_00], dim=1)\n            logits_1011 = torch.cat([logits_10, logits_11], dim=1)\n            logits = torch.cat([logits_0100, logits_1011], dim=0)\n\n            # create labels\n            labels = torch.arange(batch_size, device=device, dtype=torch.long)\n            if self.gather_distributed:\n                labels = labels + dist.rank() * batch_size\n            labels = labels.repeat(2)\n\n        loss = self.cross_entropy(logits, labels)\n\n        return loss\n", "path": "lightly/loss/ntx_ent_loss.py"}]}