We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Loss stuck
Hi, I am trying to run the tutorial posted here 
https://docs.lightly.ai/self-supervised-learning/tutorials/package/tutorial_moco_memory_bank.html
But my loss is stuck at 8.32 after 100 epochs
python 3.9
pytorch-lightning 1.8.1 
lightly 1.2.38

Any suggestions on how I should troubleshoot this?
Thanks in advance!

--- 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         self.register_buffer("bank", tensor=torch.empty(0, dtype=torch.float), persistent=False)
50         self.register_buffer("bank_ptr", tensor=torch.empty(0, dtype=torch.long), persistent=False)
51 
52     @torch.no_grad()
53     def _init_memory_bank(self, dim: int):
54         """Initialize the memory bank if it's empty
55 
56         Args:
57             dim:
58                 The dimension of the which are stored in the bank.
59 
60         """
61         # create memory bank
62         # we could use register buffers like in the moco repo
63         # https://github.com/facebookresearch/moco but we don't
64         # want to pollute our checkpoints
65         self.bank = torch.randn(dim, self.size).type_as(self.bank)
66         torch.nn.functional.normalize(self.bank, dim=0)
67         self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)
68 
69     @torch.no_grad()
70     def _dequeue_and_enqueue(self, batch: torch.Tensor):
71         """Dequeue the oldest batch and add the latest one
72 
73         Args:
74             batch:
75                 The latest batch of keys to add to the memory bank.
76 
77         """
78         batch_size = batch.shape[0]
79         ptr = int(self.bank_ptr)
80 
81         if ptr + batch_size >= self.size:
82             self.bank[:, ptr:] = batch[:self.size - ptr].T.detach()
83             self.bank_ptr[0] = 0
84         else:
85             self.bank[:, ptr:ptr + batch_size] = batch.T.detach()
86             self.bank_ptr[0] = ptr + batch_size
87 
88     def forward(self,
89                 output: torch.Tensor,
90                 labels: torch.Tensor = None,
91                 update: bool = False):
92         """Query memory bank for additional negative samples
93 
94         Args:
95             output:
96                 The output of the model.
97             labels:
98                 Should always be None, will be ignored.
99 
100         Returns:
101             The output if the memory bank is of size 0, otherwise the output
102             and the entries from the memory bank.
103 
104         """
105 
106         # no memory bank, return the output
107         if self.size == 0:
108             return output, None
109 
110         _, dim = output.shape
111 
112         # initialize the memory bank if it is not already done
113         if self.bank.nelement() == 0:
114             self._init_memory_bank(dim)
115 
116         # query and update memory bank
117         bank = self.bank.clone().detach()
118 
119         # only update memory bank if we later do backward pass (gradient)
120         if update:
121             self._dequeue_and_enqueue(output)
122 
123         return output, bank
124
```
--- 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
@@ -63,7 +63,7 @@
         # https://github.com/facebookresearch/moco but we don't
         # want to pollute our checkpoints
         self.bank = torch.randn(dim, self.size).type_as(self.bank)
-        torch.nn.functional.normalize(self.bank, dim=0)
+        self.bank = torch.nn.functional.normalize(self.bank, dim=0)
         self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)
 
     @torch.no_grad()

{"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@@ -63,7 +63,7 @@\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 = 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", "issue": "Loss stuck\nHi, I am trying to run the tutorial posted here \r\nhttps://docs.lightly.ai/self-supervised-learning/tutorials/package/tutorial_moco_memory_bank.html\r\nBut my loss is stuck at 8.32 after 100 epochs\r\npython 3.9\r\npytorch-lightning 1.8.1 \r\nlightly 1.2.38\r\n\r\nAny suggestions on how I should troubleshoot this?\r\nThanks in advance!\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        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"}], "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        self.bank = 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"}]}