We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error Using val_online_cls_top1
    print(f"max {metric}: {max(metric_callback.val_metrics[metric])}")
KeyError: 'val_online_cls_top1'

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `benchmarks/imagenet/resnet50/swav.py`
Content:
```
1 import math
2 from typing import List, Tuple
3 
4 import torch
5 from pytorch_lightning import LightningModule
6 from torch import Tensor
7 from torch.nn import Identity, ModuleList
8 from torch.nn import functional as F
9 from torchvision.models import resnet50
10 
11 from lightly.loss.memory_bank import MemoryBankModule
12 from lightly.loss.swav_loss import SwaVLoss
13 from lightly.models.modules import SwaVProjectionHead, SwaVPrototypes
14 from lightly.models.utils import get_weight_decay_parameters
15 from lightly.transforms import SwaVTransform
16 from lightly.utils.benchmarking import OnlineLinearClassifier
17 from lightly.utils.lars import LARS
18 from lightly.utils.scheduler import CosineWarmupScheduler
19 
20 CROP_COUNTS: Tuple[int, int] = (2, 6)
21 
22 
23 class SwAV(LightningModule):
24     def __init__(self, batch_size_per_device: int, num_classes: int) -> None:
25         super().__init__()
26         self.save_hyperparameters()
27         self.batch_size_per_device = batch_size_per_device
28 
29         resnet = resnet50()
30         resnet.fc = Identity()  # Ignore classification head
31         self.backbone = resnet
32         self.projection_head = SwaVProjectionHead()
33         self.prototypes = SwaVPrototypes(n_steps_frozen_prototypes=1)
34         self.criterion = SwaVLoss(sinkhorn_gather_distributed=True)
35         self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)
36 
37         # Use a queue for small batch sizes (<= 256).
38         self.start_queue_at_epoch = 15
39         self.n_batches_in_queue = 15
40         self.queues = ModuleList(
41             [
42                 MemoryBankModule(
43                     size=self.n_batches_in_queue * self.batch_size_per_device
44                 )
45                 for _ in range(CROP_COUNTS[0])
46             ]
47         )
48 
49     def forward(self, x: Tensor) -> Tensor:
50         return self.backbone(x)
51 
52     def project(self, x: Tensor) -> Tensor:
53         x = self.projection_head(x)
54         return F.normalize(x, dim=1, p=2)
55 
56     def training_step(
57         self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int
58     ) -> Tensor:
59         # Normalize the prototypes so they are on the unit sphere.
60         self.prototypes.normalize()
61 
62         # The dataloader returns a list of image crops where the
63         # first few items are high resolution crops and the rest are low
64         # resolution crops.
65         multi_crops, targets = batch[0], batch[1]
66 
67         # Forward pass through backbone and projection head.
68         multi_crop_features = [
69             self.forward(crops).flatten(start_dim=1) for crops in multi_crops
70         ]
71         multi_crop_projections = [
72             self.project(features) for features in multi_crop_features
73         ]
74 
75         # Get the queue projections and logits.
76         queue_crop_logits = None
77         with torch.no_grad():
78             if self.current_epoch >= self.start_queue_at_epoch:
79                 # Start filling the queue.
80                 queue_crop_projections = _update_queue(
81                     projections=multi_crop_projections[: CROP_COUNTS[0]],
82                     queues=self.queues,
83                 )
84                 if batch_idx > self.n_batches_in_queue:
85                     # The queue is filled, so we can start using it.
86                     queue_crop_logits = [
87                         self.prototypes(projections, step=self.current_epoch)
88                         for projections in queue_crop_projections
89                     ]
90 
91         # Get the rest of the multi-crop logits.
92         multi_crop_logits = [
93             self.prototypes(projections, step=self.current_epoch)
94             for projections in multi_crop_projections
95         ]
96 
97         # Calculate the SwAV loss.
98         loss = self.criterion(
99             high_resolution_outputs=multi_crop_logits[: CROP_COUNTS[0]],
100             low_resolution_outputs=multi_crop_logits[CROP_COUNTS[0] :],
101             queue_outputs=queue_crop_logits,
102         )
103         self.log(
104             "train_loss",
105             loss,
106             prog_bar=True,
107             sync_dist=True,
108             batch_size_per_device=len(targets),
109         )
110 
111         # Calculate the classification loss.
112         cls_loss, cls_log = self.online_classifier.training_step(
113             (multi_crop_features[0].detach(), targets), batch_idx
114         )
115         self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets))
116         return loss + cls_loss
117 
118     def validation_step(
119         self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int
120     ) -> Tensor:
121         images, targets = batch[0], batch[1]
122         features = self.forward(images).flatten(start_dim=1)
123         cls_loss, cls_log = self.online_classifier.validation_step(
124             (features.detach(), targets), batch_idx
125         )
126         self.log_dict(
127             cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)
128         )
129         return cls_loss
130 
131     def configure_optimizers(self):
132         # Don't use weight decay for batch norm, bias parameters, and classification
133         # head to improve performance.
134         params, params_no_weight_decay = get_weight_decay_parameters(
135             [self.backbone, self.projection_head, self.prototypes]
136         )
137         optimizer = LARS(
138             [
139                 {"name": "swav", "params": params},
140                 {
141                     "name": "swav_no_weight_decay",
142                     "params": params_no_weight_decay,
143                     "weight_decay": 0.0,
144                 },
145                 {
146                     "name": "online_classifier",
147                     "params": self.online_classifier.parameters(),
148                     "weight_decay": 0.0,
149                 },
150             ],
151             # Smaller learning rate for smaller batches: lr=0.6 for batch_size=256
152             # scaled linearly by batch size to lr=4.8 for batch_size=2048.
153             # See Appendix A.1. and A.6. in SwAV paper https://arxiv.org/pdf/2006.09882.pdf
154             lr=0.6 * (self.batch_size_per_device * self.trainer.world_size) / 256,
155             momentum=0.9,
156             weight_decay=1e-6,
157         )
158         scheduler = {
159             "scheduler": CosineWarmupScheduler(
160                 optimizer=optimizer,
161                 warmup_epochs=(
162                     self.trainer.estimated_stepping_batches
163                     / self.trainer.max_epochs
164                     * 10
165                 ),
166                 max_epochs=self.trainer.estimated_stepping_batches,
167                 end_value=0.0006
168                 * (self.batch_size_per_device * self.trainer.world_size)
169                 / 256,
170             ),
171             "interval": "step",
172         }
173         return [optimizer], [scheduler]
174 
175 
176 transform = SwaVTransform(crop_counts=CROP_COUNTS)
177 
178 
179 @torch.no_grad()
180 def _update_queue(
181     projections: List[Tensor],
182     queues: ModuleList,
183 ):
184     """Adds the high resolution projections to the queues and returns the queues."""
185 
186     if len(projections) != len(queues):
187         raise ValueError(
188             f"The number of queues ({len(queues)}) should be equal to the number of high "
189             f"resolution inputs ({len(projections)})."
190         )
191 
192     # Get the queue projections
193     queue_projections = []
194     for i in range(len(queues)):
195         _, projections = queues[i](projections[i], update=True)
196         # Queue projections are in (num_ftrs X queue_length) shape, while the high res
197         # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.
198         projections = torch.permute(projections, (1, 0))
199         queue_projections.append(projections)
200 
201     return queue_projections
202
```
--- 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/benchmarks/imagenet/resnet50/swav.py b/benchmarks/imagenet/resnet50/swav.py
--- a/benchmarks/imagenet/resnet50/swav.py
+++ b/benchmarks/imagenet/resnet50/swav.py
@@ -105,14 +105,14 @@
             loss,
             prog_bar=True,
             sync_dist=True,
-            batch_size_per_device=len(targets),
+            batch_size=len(targets),
         )
 
         # Calculate the classification loss.
         cls_loss, cls_log = self.online_classifier.training_step(
             (multi_crop_features[0].detach(), targets), batch_idx
         )
-        self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets))
+        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))
         return loss + cls_loss
 
     def validation_step(
@@ -123,9 +123,7 @@
         cls_loss, cls_log = self.online_classifier.validation_step(
             (features.detach(), targets), batch_idx
         )
-        self.log_dict(
-            cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)
-        )
+        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))
         return cls_loss
 
     def configure_optimizers(self):
@@ -192,10 +190,10 @@
     # Get the queue projections
     queue_projections = []
     for i in range(len(queues)):
-        _, projections = queues[i](projections[i], update=True)
+        _, queue_proj = queues[i](projections[i], update=True)
         # Queue projections are in (num_ftrs X queue_length) shape, while the high res
         # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.
-        projections = torch.permute(projections, (1, 0))
-        queue_projections.append(projections)
+        queue_proj = torch.permute(queue_proj, (1, 0))
+        queue_projections.append(queue_proj)
 
     return queue_projections

{"golden_diff": "diff --git a/benchmarks/imagenet/resnet50/swav.py b/benchmarks/imagenet/resnet50/swav.py\n--- a/benchmarks/imagenet/resnet50/swav.py\n+++ b/benchmarks/imagenet/resnet50/swav.py\n@@ -105,14 +105,14 @@\n             loss,\n             prog_bar=True,\n             sync_dist=True,\n-            batch_size_per_device=len(targets),\n+            batch_size=len(targets),\n         )\n \n         # Calculate the classification loss.\n         cls_loss, cls_log = self.online_classifier.training_step(\n             (multi_crop_features[0].detach(), targets), batch_idx\n         )\n-        self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets))\n+        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))\n         return loss + cls_loss\n \n     def validation_step(\n@@ -123,9 +123,7 @@\n         cls_loss, cls_log = self.online_classifier.validation_step(\n             (features.detach(), targets), batch_idx\n         )\n-        self.log_dict(\n-            cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)\n-        )\n+        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))\n         return cls_loss\n \n     def configure_optimizers(self):\n@@ -192,10 +190,10 @@\n     # Get the queue projections\n     queue_projections = []\n     for i in range(len(queues)):\n-        _, projections = queues[i](projections[i], update=True)\n+        _, queue_proj = queues[i](projections[i], update=True)\n         # Queue projections are in (num_ftrs X queue_length) shape, while the high res\n         # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.\n-        projections = torch.permute(projections, (1, 0))\n-        queue_projections.append(projections)\n+        queue_proj = torch.permute(queue_proj, (1, 0))\n+        queue_projections.append(queue_proj)\n \n     return queue_projections\n", "issue": "Error Using val_online_cls_top1\n    print(f\"max {metric}: {max(metric_callback.val_metrics[metric])}\")\r\nKeyError: 'val_online_cls_top1'\n", "before_files": [{"content": "import math\nfrom typing import List, Tuple\n\nimport torch\nfrom pytorch_lightning import LightningModule\nfrom torch import Tensor\nfrom torch.nn import Identity, ModuleList\nfrom torch.nn import functional as F\nfrom torchvision.models import resnet50\n\nfrom lightly.loss.memory_bank import MemoryBankModule\nfrom lightly.loss.swav_loss import SwaVLoss\nfrom lightly.models.modules import SwaVProjectionHead, SwaVPrototypes\nfrom lightly.models.utils import get_weight_decay_parameters\nfrom lightly.transforms import SwaVTransform\nfrom lightly.utils.benchmarking import OnlineLinearClassifier\nfrom lightly.utils.lars import LARS\nfrom lightly.utils.scheduler import CosineWarmupScheduler\n\nCROP_COUNTS: Tuple[int, int] = (2, 6)\n\n\nclass SwAV(LightningModule):\n    def __init__(self, batch_size_per_device: int, num_classes: int) -> None:\n        super().__init__()\n        self.save_hyperparameters()\n        self.batch_size_per_device = batch_size_per_device\n\n        resnet = resnet50()\n        resnet.fc = Identity()  # Ignore classification head\n        self.backbone = resnet\n        self.projection_head = SwaVProjectionHead()\n        self.prototypes = SwaVPrototypes(n_steps_frozen_prototypes=1)\n        self.criterion = SwaVLoss(sinkhorn_gather_distributed=True)\n        self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)\n\n        # Use a queue for small batch sizes (<= 256).\n        self.start_queue_at_epoch = 15\n        self.n_batches_in_queue = 15\n        self.queues = ModuleList(\n            [\n                MemoryBankModule(\n                    size=self.n_batches_in_queue * self.batch_size_per_device\n                )\n                for _ in range(CROP_COUNTS[0])\n            ]\n        )\n\n    def forward(self, x: Tensor) -> Tensor:\n        return self.backbone(x)\n\n    def project(self, x: Tensor) -> Tensor:\n        x = self.projection_head(x)\n        return F.normalize(x, dim=1, p=2)\n\n    def training_step(\n        self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int\n    ) -> Tensor:\n        # Normalize the prototypes so they are on the unit sphere.\n        self.prototypes.normalize()\n\n        # The dataloader returns a list of image crops where the\n        # first few items are high resolution crops and the rest are low\n        # resolution crops.\n        multi_crops, targets = batch[0], batch[1]\n\n        # Forward pass through backbone and projection head.\n        multi_crop_features = [\n            self.forward(crops).flatten(start_dim=1) for crops in multi_crops\n        ]\n        multi_crop_projections = [\n            self.project(features) for features in multi_crop_features\n        ]\n\n        # Get the queue projections and logits.\n        queue_crop_logits = None\n        with torch.no_grad():\n            if self.current_epoch >= self.start_queue_at_epoch:\n                # Start filling the queue.\n                queue_crop_projections = _update_queue(\n                    projections=multi_crop_projections[: CROP_COUNTS[0]],\n                    queues=self.queues,\n                )\n                if batch_idx > self.n_batches_in_queue:\n                    # The queue is filled, so we can start using it.\n                    queue_crop_logits = [\n                        self.prototypes(projections, step=self.current_epoch)\n                        for projections in queue_crop_projections\n                    ]\n\n        # Get the rest of the multi-crop logits.\n        multi_crop_logits = [\n            self.prototypes(projections, step=self.current_epoch)\n            for projections in multi_crop_projections\n        ]\n\n        # Calculate the SwAV loss.\n        loss = self.criterion(\n            high_resolution_outputs=multi_crop_logits[: CROP_COUNTS[0]],\n            low_resolution_outputs=multi_crop_logits[CROP_COUNTS[0] :],\n            queue_outputs=queue_crop_logits,\n        )\n        self.log(\n            \"train_loss\",\n            loss,\n            prog_bar=True,\n            sync_dist=True,\n            batch_size_per_device=len(targets),\n        )\n\n        # Calculate the classification loss.\n        cls_loss, cls_log = self.online_classifier.training_step(\n            (multi_crop_features[0].detach(), targets), batch_idx\n        )\n        self.log_dict(cls_log, sync_dist=True, batch_size_per_device=len(targets))\n        return loss + cls_loss\n\n    def validation_step(\n        self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int\n    ) -> Tensor:\n        images, targets = batch[0], batch[1]\n        features = self.forward(images).flatten(start_dim=1)\n        cls_loss, cls_log = self.online_classifier.validation_step(\n            (features.detach(), targets), batch_idx\n        )\n        self.log_dict(\n            cls_log, prog_bar=True, sync_dist=True, batch_size_per_device=len(targets)\n        )\n        return cls_loss\n\n    def configure_optimizers(self):\n        # Don't use weight decay for batch norm, bias parameters, and classification\n        # head to improve performance.\n        params, params_no_weight_decay = get_weight_decay_parameters(\n            [self.backbone, self.projection_head, self.prototypes]\n        )\n        optimizer = LARS(\n            [\n                {\"name\": \"swav\", \"params\": params},\n                {\n                    \"name\": \"swav_no_weight_decay\",\n                    \"params\": params_no_weight_decay,\n                    \"weight_decay\": 0.0,\n                },\n                {\n                    \"name\": \"online_classifier\",\n                    \"params\": self.online_classifier.parameters(),\n                    \"weight_decay\": 0.0,\n                },\n            ],\n            # Smaller learning rate for smaller batches: lr=0.6 for batch_size=256\n            # scaled linearly by batch size to lr=4.8 for batch_size=2048.\n            # See Appendix A.1. and A.6. in SwAV paper https://arxiv.org/pdf/2006.09882.pdf\n            lr=0.6 * (self.batch_size_per_device * self.trainer.world_size) / 256,\n            momentum=0.9,\n            weight_decay=1e-6,\n        )\n        scheduler = {\n            \"scheduler\": CosineWarmupScheduler(\n                optimizer=optimizer,\n                warmup_epochs=(\n                    self.trainer.estimated_stepping_batches\n                    / self.trainer.max_epochs\n                    * 10\n                ),\n                max_epochs=self.trainer.estimated_stepping_batches,\n                end_value=0.0006\n                * (self.batch_size_per_device * self.trainer.world_size)\n                / 256,\n            ),\n            \"interval\": \"step\",\n        }\n        return [optimizer], [scheduler]\n\n\ntransform = SwaVTransform(crop_counts=CROP_COUNTS)\n\n\n@torch.no_grad()\ndef _update_queue(\n    projections: List[Tensor],\n    queues: ModuleList,\n):\n    \"\"\"Adds the high resolution projections to the queues and returns the queues.\"\"\"\n\n    if len(projections) != len(queues):\n        raise ValueError(\n            f\"The number of queues ({len(queues)}) should be equal to the number of high \"\n            f\"resolution inputs ({len(projections)}).\"\n        )\n\n    # Get the queue projections\n    queue_projections = []\n    for i in range(len(queues)):\n        _, projections = queues[i](projections[i], update=True)\n        # Queue projections are in (num_ftrs X queue_length) shape, while the high res\n        # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.\n        projections = torch.permute(projections, (1, 0))\n        queue_projections.append(projections)\n\n    return queue_projections\n", "path": "benchmarks/imagenet/resnet50/swav.py"}], "after_files": [{"content": "import math\nfrom typing import List, Tuple\n\nimport torch\nfrom pytorch_lightning import LightningModule\nfrom torch import Tensor\nfrom torch.nn import Identity, ModuleList\nfrom torch.nn import functional as F\nfrom torchvision.models import resnet50\n\nfrom lightly.loss.memory_bank import MemoryBankModule\nfrom lightly.loss.swav_loss import SwaVLoss\nfrom lightly.models.modules import SwaVProjectionHead, SwaVPrototypes\nfrom lightly.models.utils import get_weight_decay_parameters\nfrom lightly.transforms import SwaVTransform\nfrom lightly.utils.benchmarking import OnlineLinearClassifier\nfrom lightly.utils.lars import LARS\nfrom lightly.utils.scheduler import CosineWarmupScheduler\n\nCROP_COUNTS: Tuple[int, int] = (2, 6)\n\n\nclass SwAV(LightningModule):\n    def __init__(self, batch_size_per_device: int, num_classes: int) -> None:\n        super().__init__()\n        self.save_hyperparameters()\n        self.batch_size_per_device = batch_size_per_device\n\n        resnet = resnet50()\n        resnet.fc = Identity()  # Ignore classification head\n        self.backbone = resnet\n        self.projection_head = SwaVProjectionHead()\n        self.prototypes = SwaVPrototypes(n_steps_frozen_prototypes=1)\n        self.criterion = SwaVLoss(sinkhorn_gather_distributed=True)\n        self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)\n\n        # Use a queue for small batch sizes (<= 256).\n        self.start_queue_at_epoch = 15\n        self.n_batches_in_queue = 15\n        self.queues = ModuleList(\n            [\n                MemoryBankModule(\n                    size=self.n_batches_in_queue * self.batch_size_per_device\n                )\n                for _ in range(CROP_COUNTS[0])\n            ]\n        )\n\n    def forward(self, x: Tensor) -> Tensor:\n        return self.backbone(x)\n\n    def project(self, x: Tensor) -> Tensor:\n        x = self.projection_head(x)\n        return F.normalize(x, dim=1, p=2)\n\n    def training_step(\n        self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int\n    ) -> Tensor:\n        # Normalize the prototypes so they are on the unit sphere.\n        self.prototypes.normalize()\n\n        # The dataloader returns a list of image crops where the\n        # first few items are high resolution crops and the rest are low\n        # resolution crops.\n        multi_crops, targets = batch[0], batch[1]\n\n        # Forward pass through backbone and projection head.\n        multi_crop_features = [\n            self.forward(crops).flatten(start_dim=1) for crops in multi_crops\n        ]\n        multi_crop_projections = [\n            self.project(features) for features in multi_crop_features\n        ]\n\n        # Get the queue projections and logits.\n        queue_crop_logits = None\n        with torch.no_grad():\n            if self.current_epoch >= self.start_queue_at_epoch:\n                # Start filling the queue.\n                queue_crop_projections = _update_queue(\n                    projections=multi_crop_projections[: CROP_COUNTS[0]],\n                    queues=self.queues,\n                )\n                if batch_idx > self.n_batches_in_queue:\n                    # The queue is filled, so we can start using it.\n                    queue_crop_logits = [\n                        self.prototypes(projections, step=self.current_epoch)\n                        for projections in queue_crop_projections\n                    ]\n\n        # Get the rest of the multi-crop logits.\n        multi_crop_logits = [\n            self.prototypes(projections, step=self.current_epoch)\n            for projections in multi_crop_projections\n        ]\n\n        # Calculate the SwAV loss.\n        loss = self.criterion(\n            high_resolution_outputs=multi_crop_logits[: CROP_COUNTS[0]],\n            low_resolution_outputs=multi_crop_logits[CROP_COUNTS[0] :],\n            queue_outputs=queue_crop_logits,\n        )\n        self.log(\n            \"train_loss\",\n            loss,\n            prog_bar=True,\n            sync_dist=True,\n            batch_size=len(targets),\n        )\n\n        # Calculate the classification loss.\n        cls_loss, cls_log = self.online_classifier.training_step(\n            (multi_crop_features[0].detach(), targets), batch_idx\n        )\n        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))\n        return loss + cls_loss\n\n    def validation_step(\n        self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int\n    ) -> Tensor:\n        images, targets = batch[0], batch[1]\n        features = self.forward(images).flatten(start_dim=1)\n        cls_loss, cls_log = self.online_classifier.validation_step(\n            (features.detach(), targets), batch_idx\n        )\n        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))\n        return cls_loss\n\n    def configure_optimizers(self):\n        # Don't use weight decay for batch norm, bias parameters, and classification\n        # head to improve performance.\n        params, params_no_weight_decay = get_weight_decay_parameters(\n            [self.backbone, self.projection_head, self.prototypes]\n        )\n        optimizer = LARS(\n            [\n                {\"name\": \"swav\", \"params\": params},\n                {\n                    \"name\": \"swav_no_weight_decay\",\n                    \"params\": params_no_weight_decay,\n                    \"weight_decay\": 0.0,\n                },\n                {\n                    \"name\": \"online_classifier\",\n                    \"params\": self.online_classifier.parameters(),\n                    \"weight_decay\": 0.0,\n                },\n            ],\n            # Smaller learning rate for smaller batches: lr=0.6 for batch_size=256\n            # scaled linearly by batch size to lr=4.8 for batch_size=2048.\n            # See Appendix A.1. and A.6. in SwAV paper https://arxiv.org/pdf/2006.09882.pdf\n            lr=0.6 * (self.batch_size_per_device * self.trainer.world_size) / 256,\n            momentum=0.9,\n            weight_decay=1e-6,\n        )\n        scheduler = {\n            \"scheduler\": CosineWarmupScheduler(\n                optimizer=optimizer,\n                warmup_epochs=(\n                    self.trainer.estimated_stepping_batches\n                    / self.trainer.max_epochs\n                    * 10\n                ),\n                max_epochs=self.trainer.estimated_stepping_batches,\n                end_value=0.0006\n                * (self.batch_size_per_device * self.trainer.world_size)\n                / 256,\n            ),\n            \"interval\": \"step\",\n        }\n        return [optimizer], [scheduler]\n\n\ntransform = SwaVTransform(crop_counts=CROP_COUNTS)\n\n\n@torch.no_grad()\ndef _update_queue(\n    projections: List[Tensor],\n    queues: ModuleList,\n):\n    \"\"\"Adds the high resolution projections to the queues and returns the queues.\"\"\"\n\n    if len(projections) != len(queues):\n        raise ValueError(\n            f\"The number of queues ({len(queues)}) should be equal to the number of high \"\n            f\"resolution inputs ({len(projections)}).\"\n        )\n\n    # Get the queue projections\n    queue_projections = []\n    for i in range(len(queues)):\n        _, queue_proj = queues[i](projections[i], update=True)\n        # Queue projections are in (num_ftrs X queue_length) shape, while the high res\n        # projections are in (batch_size_per_device X num_ftrs). Swap the axes for interoperability.\n        queue_proj = torch.permute(queue_proj, (1, 0))\n        queue_projections.append(queue_proj)\n\n    return queue_projections\n", "path": "benchmarks/imagenet/resnet50/swav.py"}]}