We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
NNMemoryBank not working with DataParallel
I have been using the NNMemoryBank as a component in my module and noticed that at each forward pass `NNMemoryBank.bank` is equal to `None`. This only occurs when my module is wrapped in `DataParallel`. As a result, throughout training my NN pairs are always random noise (surprisingly, this only hurt the contrastive learning performance by a few percentage point on linear probing??).

Here is a simple test case that highlights the issue:
```
import torch
from lightly.models.modules import NNMemoryBankModule
memory_bank = NNMemoryBankModule(size=1000)
print(memory_bank.bank)
memory_bank(torch.randn((100, 512)))
print(memory_bank.bank)

memory_bank = NNMemoryBankModule(size=1000)
memory_bank = torch.nn.DataParallel(memory_bank, device_ids=[0,1])
print(memory_bank.module.bank)
memory_bank(torch.randn((100, 512)))
print(memory_bank.module.bank)
```

The output of the first is `None` and a `torch.Tensor`, as expected. The output for the second is `None` for both.

--- 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/memory_bank.py`
Content:
```
1 """ Memory Bank Wrapper """
2 
3 # Copyright (c) 2020. Lightly AG and its affiliates.
4 # All Rights Reserved
5 
6 import torch
7 import functools
8 
9 class MemoryBankModule(torch.nn.Module):
10     """Memory bank implementation
11 
12     This is a parent class to all loss functions implemented by the lightly
13     Python package. This way, any loss can be used with a memory bank if 
14     desired.
15 
16     Attributes:
17         size:
18             Number of keys the memory bank can store. If set to 0,
19             memory bank is not used.
20 
21     Examples:
22         >>> class MyLossFunction(MemoryBankModule):
23         >>>
24         >>>     def __init__(self, memory_bank_size: int = 2 ** 16):
25         >>>         super(MyLossFunction, self).__init__(memory_bank_size)
26         >>>
27         >>>     def forward(self, output: torch.Tensor,
28         >>>                 labels: torch.Tensor = None):
29         >>>
30         >>>         output, negatives = super(
31         >>>             MyLossFunction, self).forward(output)
32         >>>
33         >>>         if negatives is not None:
34         >>>             # evaluate loss with negative samples
35         >>>         else:
36         >>>             # evaluate loss without negative samples
37 
38     """
39 
40     def __init__(self, size: int = 2 ** 16):
41 
42         super(MemoryBankModule, self).__init__()
43 
44         if size < 0:
45             msg = f'Illegal memory bank size {size}, must be non-negative.'
46             raise ValueError(msg)
47 
48         self.size = size
49 
50         self.bank = None
51         self.bank_ptr = None
52     
53     @torch.no_grad()
54     def _init_memory_bank(self, dim: int):
55         """Initialize the memory bank if it's empty
56 
57         Args:
58             dim:
59                 The dimension of the which are stored in the bank.
60 
61         """
62         # create memory bank
63         # we could use register buffers like in the moco repo
64         # https://github.com/facebookresearch/moco but we don't
65         # want to pollute our checkpoints
66         self.bank = torch.randn(dim, self.size)
67         self.bank = torch.nn.functional.normalize(self.bank, dim=0)
68         self.bank_ptr = torch.LongTensor([0])
69 
70     @torch.no_grad()
71     def _dequeue_and_enqueue(self, batch: torch.Tensor):
72         """Dequeue the oldest batch and add the latest one
73 
74         Args:
75             batch:
76                 The latest batch of keys to add to the memory bank.
77 
78         """
79         batch_size = batch.shape[0]
80         ptr = int(self.bank_ptr)
81 
82         if ptr + batch_size >= self.size:
83             self.bank[:, ptr:] = batch[:self.size - ptr].T.detach()
84             self.bank_ptr[0] = 0
85         else:
86             self.bank[:, ptr:ptr + batch_size] = batch.T.detach()
87             self.bank_ptr[0] = ptr + batch_size
88 
89     def forward(self,
90                 output: torch.Tensor,
91                 labels: torch.Tensor = None,
92                 update: bool = False):
93         """Query memory bank for additional negative samples
94 
95         Args:
96             output:
97                 The output of the model.
98             labels:
99                 Should always be None, will be ignored.
100 
101         Returns:
102             The output if the memory bank is of size 0, otherwise the output
103             and the entries from the memory bank.
104 
105         """
106 
107         # no memory bank, return the output
108         if self.size == 0:
109             return output, None
110 
111         _, dim = output.shape
112 
113         # initialize the memory bank if it is not already done
114         if self.bank is None:
115             self._init_memory_bank(dim)
116 
117         # query and update memory bank
118         bank = self.bank.clone().detach()
119 
120         # only update memory bank if we later do backward pass (gradient)
121         if update:
122             self._dequeue_and_enqueue(output)
123 
124         return output, bank
125
```
--- 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/memory_bank.py b/lightly/loss/memory_bank.py
--- a/lightly/loss/memory_bank.py
+++ b/lightly/loss/memory_bank.py
@@ -46,10 +46,9 @@
             raise ValueError(msg)
 
         self.size = size
+        self.register_buffer("bank", tensor=torch.empty(0, dtype=torch.float), persistent=False)
+        self.register_buffer("bank_ptr", tensor=torch.empty(0, dtype=torch.long), persistent=False)
 
-        self.bank = None
-        self.bank_ptr = None
-    
     @torch.no_grad()
     def _init_memory_bank(self, dim: int):
         """Initialize the memory bank if it's empty
@@ -63,9 +62,9 @@
         # we could use register buffers like in the moco repo
         # https://github.com/facebookresearch/moco but we don't
         # want to pollute our checkpoints
-        self.bank = torch.randn(dim, self.size)
-        self.bank = torch.nn.functional.normalize(self.bank, dim=0)
-        self.bank_ptr = torch.LongTensor([0])
+        self.bank = torch.randn(dim, self.size).type_as(self.bank)
+        torch.nn.functional.normalize(self.bank, dim=0)
+        self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)
 
     @torch.no_grad()
     def _dequeue_and_enqueue(self, batch: torch.Tensor):
@@ -111,7 +110,7 @@
         _, dim = output.shape
 
         # initialize the memory bank if it is not already done
-        if self.bank is None:
+        if self.bank.nelement() == 0:
             self._init_memory_bank(dim)
 
         # query and update memory bank

{"golden_diff": "diff --git a/lightly/loss/memory_bank.py b/lightly/loss/memory_bank.py\n--- a/lightly/loss/memory_bank.py\n+++ b/lightly/loss/memory_bank.py\n@@ -46,10 +46,9 @@\n             raise ValueError(msg)\n \n         self.size = size\n+        self.register_buffer(\"bank\", tensor=torch.empty(0, dtype=torch.float), persistent=False)\n+        self.register_buffer(\"bank_ptr\", tensor=torch.empty(0, dtype=torch.long), persistent=False)\n \n-        self.bank = None\n-        self.bank_ptr = None\n-    \n     @torch.no_grad()\n     def _init_memory_bank(self, dim: int):\n         \"\"\"Initialize the memory bank if it's empty\n@@ -63,9 +62,9 @@\n         # we could use register buffers like in the moco repo\n         # https://github.com/facebookresearch/moco but we don't\n         # want to pollute our checkpoints\n-        self.bank = torch.randn(dim, self.size)\n-        self.bank = torch.nn.functional.normalize(self.bank, dim=0)\n-        self.bank_ptr = torch.LongTensor([0])\n+        self.bank = torch.randn(dim, self.size).type_as(self.bank)\n+        torch.nn.functional.normalize(self.bank, dim=0)\n+        self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)\n \n     @torch.no_grad()\n     def _dequeue_and_enqueue(self, batch: torch.Tensor):\n@@ -111,7 +110,7 @@\n         _, dim = output.shape\n \n         # initialize the memory bank if it is not already done\n-        if self.bank is None:\n+        if self.bank.nelement() == 0:\n             self._init_memory_bank(dim)\n \n         # query and update memory bank\n", "issue": "NNMemoryBank not working with DataParallel\nI have been using the NNMemoryBank as a component in my module and noticed that at each forward pass `NNMemoryBank.bank` is equal to `None`. This only occurs when my module is wrapped in `DataParallel`. As a result, throughout training my NN pairs are always random noise (surprisingly, this only hurt the contrastive learning performance by a few percentage point on linear probing??).\r\n\r\nHere is a simple test case that highlights the issue:\r\n```\r\nimport torch\r\nfrom lightly.models.modules import NNMemoryBankModule\r\nmemory_bank = NNMemoryBankModule(size=1000)\r\nprint(memory_bank.bank)\r\nmemory_bank(torch.randn((100, 512)))\r\nprint(memory_bank.bank)\r\n\r\nmemory_bank = NNMemoryBankModule(size=1000)\r\nmemory_bank = torch.nn.DataParallel(memory_bank, device_ids=[0,1])\r\nprint(memory_bank.module.bank)\r\nmemory_bank(torch.randn((100, 512)))\r\nprint(memory_bank.module.bank)\r\n```\r\n\r\nThe output of the first is `None` and a `torch.Tensor`, as expected. The output for the second is `None` for both.\n", "before_files": [{"content": "\"\"\" Memory Bank Wrapper \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nimport functools\n\nclass MemoryBankModule(torch.nn.Module):\n    \"\"\"Memory bank implementation\n\n    This is a parent class to all loss functions implemented by the lightly\n    Python package. This way, any loss can be used with a memory bank if \n    desired.\n\n    Attributes:\n        size:\n            Number of keys the memory bank can store. If set to 0,\n            memory bank is not used.\n\n    Examples:\n        >>> class MyLossFunction(MemoryBankModule):\n        >>>\n        >>>     def __init__(self, memory_bank_size: int = 2 ** 16):\n        >>>         super(MyLossFunction, self).__init__(memory_bank_size)\n        >>>\n        >>>     def forward(self, output: torch.Tensor,\n        >>>                 labels: torch.Tensor = None):\n        >>>\n        >>>         output, negatives = super(\n        >>>             MyLossFunction, self).forward(output)\n        >>>\n        >>>         if negatives is not None:\n        >>>             # evaluate loss with negative samples\n        >>>         else:\n        >>>             # evaluate loss without negative samples\n\n    \"\"\"\n\n    def __init__(self, size: int = 2 ** 16):\n\n        super(MemoryBankModule, self).__init__()\n\n        if size < 0:\n            msg = f'Illegal memory bank size {size}, must be non-negative.'\n            raise ValueError(msg)\n\n        self.size = size\n\n        self.bank = None\n        self.bank_ptr = None\n    \n    @torch.no_grad()\n    def _init_memory_bank(self, dim: int):\n        \"\"\"Initialize the memory bank if it's empty\n\n        Args:\n            dim:\n                The dimension of the which are stored in the bank.\n\n        \"\"\"\n        # create memory bank\n        # we could use register buffers like in the moco repo\n        # https://github.com/facebookresearch/moco but we don't\n        # want to pollute our checkpoints\n        self.bank = torch.randn(dim, self.size)\n        self.bank = torch.nn.functional.normalize(self.bank, dim=0)\n        self.bank_ptr = torch.LongTensor([0])\n\n    @torch.no_grad()\n    def _dequeue_and_enqueue(self, batch: torch.Tensor):\n        \"\"\"Dequeue the oldest batch and add the latest one\n\n        Args:\n            batch:\n                The latest batch of keys to add to the memory bank.\n\n        \"\"\"\n        batch_size = batch.shape[0]\n        ptr = int(self.bank_ptr)\n\n        if ptr + batch_size >= self.size:\n            self.bank[:, ptr:] = batch[:self.size - ptr].T.detach()\n            self.bank_ptr[0] = 0\n        else:\n            self.bank[:, ptr:ptr + batch_size] = batch.T.detach()\n            self.bank_ptr[0] = ptr + batch_size\n\n    def forward(self,\n                output: torch.Tensor,\n                labels: torch.Tensor = None,\n                update: bool = False):\n        \"\"\"Query memory bank for additional negative samples\n\n        Args:\n            output:\n                The output of the model.\n            labels:\n                Should always be None, will be ignored.\n\n        Returns:\n            The output if the memory bank is of size 0, otherwise the output\n            and the entries from the memory bank.\n\n        \"\"\"\n\n        # no memory bank, return the output\n        if self.size == 0:\n            return output, None\n\n        _, dim = output.shape\n\n        # initialize the memory bank if it is not already done\n        if self.bank is None:\n            self._init_memory_bank(dim)\n\n        # query and update memory bank\n        bank = self.bank.clone().detach()\n\n        # only update memory bank if we later do backward pass (gradient)\n        if update:\n            self._dequeue_and_enqueue(output)\n\n        return output, bank\n", "path": "lightly/loss/memory_bank.py"}], "after_files": [{"content": "\"\"\" Memory Bank Wrapper \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nimport functools\n\nclass MemoryBankModule(torch.nn.Module):\n    \"\"\"Memory bank implementation\n\n    This is a parent class to all loss functions implemented by the lightly\n    Python package. This way, any loss can be used with a memory bank if \n    desired.\n\n    Attributes:\n        size:\n            Number of keys the memory bank can store. If set to 0,\n            memory bank is not used.\n\n    Examples:\n        >>> class MyLossFunction(MemoryBankModule):\n        >>>\n        >>>     def __init__(self, memory_bank_size: int = 2 ** 16):\n        >>>         super(MyLossFunction, self).__init__(memory_bank_size)\n        >>>\n        >>>     def forward(self, output: torch.Tensor,\n        >>>                 labels: torch.Tensor = None):\n        >>>\n        >>>         output, negatives = super(\n        >>>             MyLossFunction, self).forward(output)\n        >>>\n        >>>         if negatives is not None:\n        >>>             # evaluate loss with negative samples\n        >>>         else:\n        >>>             # evaluate loss without negative samples\n\n    \"\"\"\n\n    def __init__(self, size: int = 2 ** 16):\n\n        super(MemoryBankModule, self).__init__()\n\n        if size < 0:\n            msg = f'Illegal memory bank size {size}, must be non-negative.'\n            raise ValueError(msg)\n\n        self.size = size\n        self.register_buffer(\"bank\", tensor=torch.empty(0, dtype=torch.float), persistent=False)\n        self.register_buffer(\"bank_ptr\", tensor=torch.empty(0, dtype=torch.long), persistent=False)\n\n    @torch.no_grad()\n    def _init_memory_bank(self, dim: int):\n        \"\"\"Initialize the memory bank if it's empty\n\n        Args:\n            dim:\n                The dimension of the which are stored in the bank.\n\n        \"\"\"\n        # create memory bank\n        # we could use register buffers like in the moco repo\n        # https://github.com/facebookresearch/moco but we don't\n        # want to pollute our checkpoints\n        self.bank = torch.randn(dim, self.size).type_as(self.bank)\n        torch.nn.functional.normalize(self.bank, dim=0)\n        self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)\n\n    @torch.no_grad()\n    def _dequeue_and_enqueue(self, batch: torch.Tensor):\n        \"\"\"Dequeue the oldest batch and add the latest one\n\n        Args:\n            batch:\n                The latest batch of keys to add to the memory bank.\n\n        \"\"\"\n        batch_size = batch.shape[0]\n        ptr = int(self.bank_ptr)\n\n        if ptr + batch_size >= self.size:\n            self.bank[:, ptr:] = batch[:self.size - ptr].T.detach()\n            self.bank_ptr[0] = 0\n        else:\n            self.bank[:, ptr:ptr + batch_size] = batch.T.detach()\n            self.bank_ptr[0] = ptr + batch_size\n\n    def forward(self,\n                output: torch.Tensor,\n                labels: torch.Tensor = None,\n                update: bool = False):\n        \"\"\"Query memory bank for additional negative samples\n\n        Args:\n            output:\n                The output of the model.\n            labels:\n                Should always be None, will be ignored.\n\n        Returns:\n            The output if the memory bank is of size 0, otherwise the output\n            and the entries from the memory bank.\n\n        \"\"\"\n\n        # no memory bank, return the output\n        if self.size == 0:\n            return output, None\n\n        _, dim = output.shape\n\n        # initialize the memory bank if it is not already done\n        if self.bank.nelement() == 0:\n            self._init_memory_bank(dim)\n\n        # query and update memory bank\n        bank = self.bank.clone().detach()\n\n        # only update memory bank if we later do backward pass (gradient)\n        if update:\n            self._dequeue_and_enqueue(output)\n\n        return output, bank\n", "path": "lightly/loss/memory_bank.py"}]}