We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
NNCLR
Hi,thanks for your work,
In NNCLR's model:
 "
# forward pass of first input x0
        f0 = self.backbone(x0).flatten(start_dim=1)
        z0 = self.projection_mlp(f0)
        p0 = self.prediction_mlp(f0)
 "

I think it's wrong, it should be that :
"
  # forward pass of first input x0
        f0 = self.backbone(x0).flatten(start_dim=1)
        z0 = self.projection_mlp(f0)
        p0 = self.prediction_mlp(z0)
" .

--- 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/models/nnclr.py`
Content:
```
1 """ NNCLR Model """
2 
3 # Copyright (c) 2021. Lightly AG and its affiliates.
4 # All Rights Reserved
5 
6 import warnings
7 
8 import torch
9 import torch.nn as nn
10 
11 from lightly.models.modules import NNCLRProjectionHead
12 from lightly.models.modules import NNCLRPredictionHead
13 
14 
15 def _prediction_mlp(in_dims: int, 
16                     h_dims: int, 
17                     out_dims: int) -> nn.Sequential:
18     """Prediction MLP. The original paper's implementation has 2 layers, with 
19     BN applied to its hidden fc layers but no BN or ReLU on the output fc layer.
20 
21     Note that the hidden dimensions should be smaller than the input/output 
22     dimensions (bottleneck structure). The default implementation using a 
23     ResNet50 backbone has an input dimension of 2048, hidden dimension of 512, 
24     and output dimension of 2048
25 
26     Args:
27         in_dims:
28             Input dimension of the first linear layer.
29         h_dims: 
30             Hidden dimension of all the fully connected layers (should be a
31             bottleneck!)
32         out_dims: 
33             Output Dimension of the final linear layer.
34 
35     Returns:
36         nn.Sequential:
37             The projection head.
38     """
39     l1 = nn.Sequential(nn.Linear(in_dims, h_dims),
40                        nn.BatchNorm1d(h_dims),
41                        nn.ReLU(inplace=True))
42 
43     l2 = nn.Linear(h_dims, out_dims)
44 
45     prediction = nn.Sequential(l1, l2)
46     return prediction
47 
48 
49 def _projection_mlp(num_ftrs: int,
50                     h_dims: int, 
51                     out_dim: int, 
52                     num_layers: int = 3) -> nn.Sequential:
53     """Projection MLP. The original paper's implementation has 3 layers, with 
54     BN applied to its hidden fc layers but no ReLU on the output fc layer. 
55     The CIFAR-10 study used a MLP with only two layers.
56 
57     Args:
58         in_dims:
59             Input dimension of the first linear layer.
60         h_dims: 
61             Hidden dimension of all the fully connected layers.
62         out_dims: 
63             Output Dimension of the final linear layer.
64         num_layers:
65             Controls the number of layers; must be 2 or 3. Defaults to 3.
66 
67     Returns:
68         nn.Sequential:
69             The projection head.
70     """
71     l1 = nn.Sequential(nn.Linear(num_ftrs, h_dims),
72                        nn.BatchNorm1d(h_dims),
73                        nn.ReLU(inplace=True))
74 
75     l2 = nn.Sequential(nn.Linear(h_dims, h_dims),
76                        nn.BatchNorm1d(h_dims),
77                        nn.ReLU(inplace=True))
78 
79     l3 = nn.Sequential(nn.Linear(h_dims, out_dim),
80                        nn.BatchNorm1d(out_dim))
81 
82     if num_layers == 3:
83         projection = nn.Sequential(l1, l2, l3)
84     elif num_layers == 2:
85         projection = nn.Sequential(l1, l3)
86     else:
87         raise NotImplementedError("Only MLPs with 2 and 3 layers are implemented.")
88 
89     return projection
90 
91 
92 class NNCLR(nn.Module):
93     """Implementation of the NNCLR[0] architecture
94 
95     Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss`
96     Recommended module: :py:class:`lightly.models.modules.nn_memory_bank.NNmemoryBankModule`
97 
98     [0] NNCLR, 2021, https://arxiv.org/abs/2104.14548
99 
100     Attributes:
101         backbone:
102             Backbone model to extract features from images.
103         num_ftrs:
104             Dimension of the embedding (before the projection head).
105         proj_hidden_dim: 
106             Dimension of the hidden layer of the projection head.
107         pred_hidden_dim:
108             Dimension of the hidden layer of the predicion head.
109         out_dim:
110             Dimension of the output (after the projection head).
111         num_mlp_layers:
112             Number of linear layers for MLP.
113 
114     Examples:
115         >>> model = NNCLR(backbone)
116         >>> criterion = NTXentLoss(temperature=0.1)
117         >>> 
118         >>> nn_replacer = NNmemoryBankModule(size=2 ** 16)
119         >>>
120         >>> # forward pass
121         >>> (z0, p0), (z1, p1) = model(x0, x1)
122         >>> z0 = nn_replacer(z0.detach(), update=False)
123         >>> z1 = nn_replacer(z1.detach(), update=True)
124         >>>
125         >>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))
126 
127     """
128 
129     def __init__(self,
130                  backbone: nn.Module,
131                  num_ftrs: int = 512,
132                  proj_hidden_dim: int = 2048,
133                  pred_hidden_dim: int = 4096,
134                  out_dim: int = 256):
135 
136         super(NNCLR, self).__init__()
137 
138         self.backbone = backbone
139         self.num_ftrs = num_ftrs
140         self.proj_hidden_dim = proj_hidden_dim
141         self.pred_hidden_dim = pred_hidden_dim
142         self.out_dim = out_dim
143 
144         self.projection_mlp = NNCLRProjectionHead(
145             num_ftrs,
146             proj_hidden_dim,
147             out_dim,
148         )
149         
150         self.prediction_mlp = NNCLRPredictionHead(
151             num_ftrs,
152             pred_hidden_dim,
153             out_dim,
154         )
155 
156         warnings.warn(Warning(
157             'The high-level building block NNCLR will be deprecated in version 1.2.0. '
158             + 'Use low-level building blocks instead. '
159             + 'See https://docs.lightly.ai/lightly.models.html for more information'),
160             PendingDeprecationWarning)
161 
162     def forward(self,
163                 x0: torch.Tensor,
164                 x1: torch.Tensor = None,
165                 return_features: bool = False):
166         """Embeds and projects the input images.
167 
168         Extracts features with the backbone and applies the projection
169         head to the output space. If both x0 and x1 are not None, both will be
170         passed through the backbone and projection head. If x1 is None, only
171         x0 will be forwarded.
172 
173         Args:
174             x0:
175                 Tensor of shape bsz x channels x W x H.
176             x1:
177                 Tensor of shape bsz x channels x W x H.
178             return_features:
179                 Whether or not to return the intermediate features backbone(x).
180 
181         Returns:
182             The output projection of x0 and (if x1 is not None) the output
183             projection of x1. If return_features is True, the output for each x
184             is a tuple (out, f) where f are the features before the projection
185             head.
186 
187         Examples:
188             >>> # single input, single output
189             >>> out = model(x) 
190             >>> 
191             >>> # single input with return_features=True
192             >>> out, f = model(x, return_features=True)
193             >>>
194             >>> # two inputs, two outputs
195             >>> out0, out1 = model(x0, x1)
196             >>>
197             >>> # two inputs, two outputs with return_features=True
198             >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True)
199 
200         """
201         
202         # forward pass of first input x0
203         f0 = self.backbone(x0).flatten(start_dim=1)
204         z0 = self.projection_mlp(f0)
205         p0 = self.prediction_mlp(f0)
206 
207         out0 = (z0, p0)
208 
209         # append features if requested
210         if return_features:
211             out0 = (out0, f0)
212 
213         # return out0 if x1 is None
214         if x1 is None:
215             return out0
216 
217         # forward pass of second input x1
218         f1 = self.backbone(x1).flatten(start_dim=1)
219         z1 = self.projection_mlp(f1)
220         p1 = self.prediction_mlp(f1)
221 
222         out1 = (z1, p1)
223 
224         # append features if requested
225         if return_features:
226             out1 = (out1, f1)
227 
228         # return both outputs
229         return out0, out1
230
```
--- 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/models/nnclr.py b/lightly/models/nnclr.py
--- a/lightly/models/nnclr.py
+++ b/lightly/models/nnclr.py
@@ -148,7 +148,7 @@
         )
         
         self.prediction_mlp = NNCLRPredictionHead(
-            num_ftrs,
+            out_dim,
             pred_hidden_dim,
             out_dim,
         )
@@ -202,7 +202,7 @@
         # forward pass of first input x0
         f0 = self.backbone(x0).flatten(start_dim=1)
         z0 = self.projection_mlp(f0)
-        p0 = self.prediction_mlp(f0)
+        p0 = self.prediction_mlp(z0)
 
         out0 = (z0, p0)
 
@@ -217,7 +217,7 @@
         # forward pass of second input x1
         f1 = self.backbone(x1).flatten(start_dim=1)
         z1 = self.projection_mlp(f1)
-        p1 = self.prediction_mlp(f1)
+        p1 = self.prediction_mlp(z1)
 
         out1 = (z1, p1)

{"golden_diff": "diff --git a/lightly/models/nnclr.py b/lightly/models/nnclr.py\n--- a/lightly/models/nnclr.py\n+++ b/lightly/models/nnclr.py\n@@ -148,7 +148,7 @@\n         )\n         \n         self.prediction_mlp = NNCLRPredictionHead(\n-            num_ftrs,\n+            out_dim,\n             pred_hidden_dim,\n             out_dim,\n         )\n@@ -202,7 +202,7 @@\n         # forward pass of first input x0\n         f0 = self.backbone(x0).flatten(start_dim=1)\n         z0 = self.projection_mlp(f0)\n-        p0 = self.prediction_mlp(f0)\n+        p0 = self.prediction_mlp(z0)\n \n         out0 = (z0, p0)\n \n@@ -217,7 +217,7 @@\n         # forward pass of second input x1\n         f1 = self.backbone(x1).flatten(start_dim=1)\n         z1 = self.projection_mlp(f1)\n-        p1 = self.prediction_mlp(f1)\n+        p1 = self.prediction_mlp(z1)\n \n         out1 = (z1, p1)\n", "issue": "NNCLR\nHi,thanks for your work,\r\nIn NNCLR's model:\r\n \"\r\n# forward pass of first input x0\r\n        f0 = self.backbone(x0).flatten(start_dim=1)\r\n        z0 = self.projection_mlp(f0)\r\n        p0 = self.prediction_mlp(f0)\r\n \"\r\n\r\nI think it's wrong, it should be that :\r\n\"\r\n  # forward pass of first input x0\r\n        f0 = self.backbone(x0).flatten(start_dim=1)\r\n        z0 = self.projection_mlp(f0)\r\n        p0 = self.prediction_mlp(z0)\r\n\" .\n", "before_files": [{"content": "\"\"\" NNCLR Model \"\"\"\n\n# Copyright (c) 2021. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport warnings\n\nimport torch\nimport torch.nn as nn\n\nfrom lightly.models.modules import NNCLRProjectionHead\nfrom lightly.models.modules import NNCLRPredictionHead\n\n\ndef _prediction_mlp(in_dims: int, \n                    h_dims: int, \n                    out_dims: int) -> nn.Sequential:\n    \"\"\"Prediction MLP. The original paper's implementation has 2 layers, with \n    BN applied to its hidden fc layers but no BN or ReLU on the output fc layer.\n\n    Note that the hidden dimensions should be smaller than the input/output \n    dimensions (bottleneck structure). The default implementation using a \n    ResNet50 backbone has an input dimension of 2048, hidden dimension of 512, \n    and output dimension of 2048\n\n    Args:\n        in_dims:\n            Input dimension of the first linear layer.\n        h_dims: \n            Hidden dimension of all the fully connected layers (should be a\n            bottleneck!)\n        out_dims: \n            Output Dimension of the final linear layer.\n\n    Returns:\n        nn.Sequential:\n            The projection head.\n    \"\"\"\n    l1 = nn.Sequential(nn.Linear(in_dims, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l2 = nn.Linear(h_dims, out_dims)\n\n    prediction = nn.Sequential(l1, l2)\n    return prediction\n\n\ndef _projection_mlp(num_ftrs: int,\n                    h_dims: int, \n                    out_dim: int, \n                    num_layers: int = 3) -> nn.Sequential:\n    \"\"\"Projection MLP. The original paper's implementation has 3 layers, with \n    BN applied to its hidden fc layers but no ReLU on the output fc layer. \n    The CIFAR-10 study used a MLP with only two layers.\n\n    Args:\n        in_dims:\n            Input dimension of the first linear layer.\n        h_dims: \n            Hidden dimension of all the fully connected layers.\n        out_dims: \n            Output Dimension of the final linear layer.\n        num_layers:\n            Controls the number of layers; must be 2 or 3. Defaults to 3.\n\n    Returns:\n        nn.Sequential:\n            The projection head.\n    \"\"\"\n    l1 = nn.Sequential(nn.Linear(num_ftrs, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l2 = nn.Sequential(nn.Linear(h_dims, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l3 = nn.Sequential(nn.Linear(h_dims, out_dim),\n                       nn.BatchNorm1d(out_dim))\n\n    if num_layers == 3:\n        projection = nn.Sequential(l1, l2, l3)\n    elif num_layers == 2:\n        projection = nn.Sequential(l1, l3)\n    else:\n        raise NotImplementedError(\"Only MLPs with 2 and 3 layers are implemented.\")\n\n    return projection\n\n\nclass NNCLR(nn.Module):\n    \"\"\"Implementation of the NNCLR[0] architecture\n\n    Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss`\n    Recommended module: :py:class:`lightly.models.modules.nn_memory_bank.NNmemoryBankModule`\n\n    [0] NNCLR, 2021, https://arxiv.org/abs/2104.14548\n\n    Attributes:\n        backbone:\n            Backbone model to extract features from images.\n        num_ftrs:\n            Dimension of the embedding (before the projection head).\n        proj_hidden_dim: \n            Dimension of the hidden layer of the projection head.\n        pred_hidden_dim:\n            Dimension of the hidden layer of the predicion head.\n        out_dim:\n            Dimension of the output (after the projection head).\n        num_mlp_layers:\n            Number of linear layers for MLP.\n\n    Examples:\n        >>> model = NNCLR(backbone)\n        >>> criterion = NTXentLoss(temperature=0.1)\n        >>> \n        >>> nn_replacer = NNmemoryBankModule(size=2 ** 16)\n        >>>\n        >>> # forward pass\n        >>> (z0, p0), (z1, p1) = model(x0, x1)\n        >>> z0 = nn_replacer(z0.detach(), update=False)\n        >>> z1 = nn_replacer(z1.detach(), update=True)\n        >>>\n        >>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))\n\n    \"\"\"\n\n    def __init__(self,\n                 backbone: nn.Module,\n                 num_ftrs: int = 512,\n                 proj_hidden_dim: int = 2048,\n                 pred_hidden_dim: int = 4096,\n                 out_dim: int = 256):\n\n        super(NNCLR, self).__init__()\n\n        self.backbone = backbone\n        self.num_ftrs = num_ftrs\n        self.proj_hidden_dim = proj_hidden_dim\n        self.pred_hidden_dim = pred_hidden_dim\n        self.out_dim = out_dim\n\n        self.projection_mlp = NNCLRProjectionHead(\n            num_ftrs,\n            proj_hidden_dim,\n            out_dim,\n        )\n        \n        self.prediction_mlp = NNCLRPredictionHead(\n            num_ftrs,\n            pred_hidden_dim,\n            out_dim,\n        )\n\n        warnings.warn(Warning(\n            'The high-level building block NNCLR will be deprecated in version 1.2.0. '\n            + 'Use low-level building blocks instead. '\n            + 'See https://docs.lightly.ai/lightly.models.html for more information'),\n            PendingDeprecationWarning)\n\n    def forward(self,\n                x0: torch.Tensor,\n                x1: torch.Tensor = None,\n                return_features: bool = False):\n        \"\"\"Embeds and projects the input images.\n\n        Extracts features with the backbone and applies the projection\n        head to the output space. If both x0 and x1 are not None, both will be\n        passed through the backbone and projection head. If x1 is None, only\n        x0 will be forwarded.\n\n        Args:\n            x0:\n                Tensor of shape bsz x channels x W x H.\n            x1:\n                Tensor of shape bsz x channels x W x H.\n            return_features:\n                Whether or not to return the intermediate features backbone(x).\n\n        Returns:\n            The output projection of x0 and (if x1 is not None) the output\n            projection of x1. If return_features is True, the output for each x\n            is a tuple (out, f) where f are the features before the projection\n            head.\n\n        Examples:\n            >>> # single input, single output\n            >>> out = model(x) \n            >>> \n            >>> # single input with return_features=True\n            >>> out, f = model(x, return_features=True)\n            >>>\n            >>> # two inputs, two outputs\n            >>> out0, out1 = model(x0, x1)\n            >>>\n            >>> # two inputs, two outputs with return_features=True\n            >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True)\n\n        \"\"\"\n        \n        # forward pass of first input x0\n        f0 = self.backbone(x0).flatten(start_dim=1)\n        z0 = self.projection_mlp(f0)\n        p0 = self.prediction_mlp(f0)\n\n        out0 = (z0, p0)\n\n        # append features if requested\n        if return_features:\n            out0 = (out0, f0)\n\n        # return out0 if x1 is None\n        if x1 is None:\n            return out0\n\n        # forward pass of second input x1\n        f1 = self.backbone(x1).flatten(start_dim=1)\n        z1 = self.projection_mlp(f1)\n        p1 = self.prediction_mlp(f1)\n\n        out1 = (z1, p1)\n\n        # append features if requested\n        if return_features:\n            out1 = (out1, f1)\n\n        # return both outputs\n        return out0, out1\n", "path": "lightly/models/nnclr.py"}], "after_files": [{"content": "\"\"\" NNCLR Model \"\"\"\n\n# Copyright (c) 2021. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport warnings\n\nimport torch\nimport torch.nn as nn\n\nfrom lightly.models.modules import NNCLRProjectionHead\nfrom lightly.models.modules import NNCLRPredictionHead\n\n\ndef _prediction_mlp(in_dims: int, \n                    h_dims: int, \n                    out_dims: int) -> nn.Sequential:\n    \"\"\"Prediction MLP. The original paper's implementation has 2 layers, with \n    BN applied to its hidden fc layers but no BN or ReLU on the output fc layer.\n\n    Note that the hidden dimensions should be smaller than the input/output \n    dimensions (bottleneck structure). The default implementation using a \n    ResNet50 backbone has an input dimension of 2048, hidden dimension of 512, \n    and output dimension of 2048\n\n    Args:\n        in_dims:\n            Input dimension of the first linear layer.\n        h_dims: \n            Hidden dimension of all the fully connected layers (should be a\n            bottleneck!)\n        out_dims: \n            Output Dimension of the final linear layer.\n\n    Returns:\n        nn.Sequential:\n            The projection head.\n    \"\"\"\n    l1 = nn.Sequential(nn.Linear(in_dims, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l2 = nn.Linear(h_dims, out_dims)\n\n    prediction = nn.Sequential(l1, l2)\n    return prediction\n\n\ndef _projection_mlp(num_ftrs: int,\n                    h_dims: int, \n                    out_dim: int, \n                    num_layers: int = 3) -> nn.Sequential:\n    \"\"\"Projection MLP. The original paper's implementation has 3 layers, with \n    BN applied to its hidden fc layers but no ReLU on the output fc layer. \n    The CIFAR-10 study used a MLP with only two layers.\n\n    Args:\n        in_dims:\n            Input dimension of the first linear layer.\n        h_dims: \n            Hidden dimension of all the fully connected layers.\n        out_dims: \n            Output Dimension of the final linear layer.\n        num_layers:\n            Controls the number of layers; must be 2 or 3. Defaults to 3.\n\n    Returns:\n        nn.Sequential:\n            The projection head.\n    \"\"\"\n    l1 = nn.Sequential(nn.Linear(num_ftrs, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l2 = nn.Sequential(nn.Linear(h_dims, h_dims),\n                       nn.BatchNorm1d(h_dims),\n                       nn.ReLU(inplace=True))\n\n    l3 = nn.Sequential(nn.Linear(h_dims, out_dim),\n                       nn.BatchNorm1d(out_dim))\n\n    if num_layers == 3:\n        projection = nn.Sequential(l1, l2, l3)\n    elif num_layers == 2:\n        projection = nn.Sequential(l1, l3)\n    else:\n        raise NotImplementedError(\"Only MLPs with 2 and 3 layers are implemented.\")\n\n    return projection\n\n\nclass NNCLR(nn.Module):\n    \"\"\"Implementation of the NNCLR[0] architecture\n\n    Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss`\n    Recommended module: :py:class:`lightly.models.modules.nn_memory_bank.NNmemoryBankModule`\n\n    [0] NNCLR, 2021, https://arxiv.org/abs/2104.14548\n\n    Attributes:\n        backbone:\n            Backbone model to extract features from images.\n        num_ftrs:\n            Dimension of the embedding (before the projection head).\n        proj_hidden_dim: \n            Dimension of the hidden layer of the projection head.\n        pred_hidden_dim:\n            Dimension of the hidden layer of the predicion head.\n        out_dim:\n            Dimension of the output (after the projection head).\n        num_mlp_layers:\n            Number of linear layers for MLP.\n\n    Examples:\n        >>> model = NNCLR(backbone)\n        >>> criterion = NTXentLoss(temperature=0.1)\n        >>> \n        >>> nn_replacer = NNmemoryBankModule(size=2 ** 16)\n        >>>\n        >>> # forward pass\n        >>> (z0, p0), (z1, p1) = model(x0, x1)\n        >>> z0 = nn_replacer(z0.detach(), update=False)\n        >>> z1 = nn_replacer(z1.detach(), update=True)\n        >>>\n        >>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))\n\n    \"\"\"\n\n    def __init__(self,\n                 backbone: nn.Module,\n                 num_ftrs: int = 512,\n                 proj_hidden_dim: int = 2048,\n                 pred_hidden_dim: int = 4096,\n                 out_dim: int = 256):\n\n        super(NNCLR, self).__init__()\n\n        self.backbone = backbone\n        self.num_ftrs = num_ftrs\n        self.proj_hidden_dim = proj_hidden_dim\n        self.pred_hidden_dim = pred_hidden_dim\n        self.out_dim = out_dim\n\n        self.projection_mlp = NNCLRProjectionHead(\n            num_ftrs,\n            proj_hidden_dim,\n            out_dim,\n        )\n        \n        self.prediction_mlp = NNCLRPredictionHead(\n            out_dim,\n            pred_hidden_dim,\n            out_dim,\n        )\n\n        warnings.warn(Warning(\n            'The high-level building block NNCLR will be deprecated in version 1.2.0. '\n            + 'Use low-level building blocks instead. '\n            + 'See https://docs.lightly.ai/lightly.models.html for more information'),\n            PendingDeprecationWarning)\n\n    def forward(self,\n                x0: torch.Tensor,\n                x1: torch.Tensor = None,\n                return_features: bool = False):\n        \"\"\"Embeds and projects the input images.\n\n        Extracts features with the backbone and applies the projection\n        head to the output space. If both x0 and x1 are not None, both will be\n        passed through the backbone and projection head. If x1 is None, only\n        x0 will be forwarded.\n\n        Args:\n            x0:\n                Tensor of shape bsz x channels x W x H.\n            x1:\n                Tensor of shape bsz x channels x W x H.\n            return_features:\n                Whether or not to return the intermediate features backbone(x).\n\n        Returns:\n            The output projection of x0 and (if x1 is not None) the output\n            projection of x1. If return_features is True, the output for each x\n            is a tuple (out, f) where f are the features before the projection\n            head.\n\n        Examples:\n            >>> # single input, single output\n            >>> out = model(x) \n            >>> \n            >>> # single input with return_features=True\n            >>> out, f = model(x, return_features=True)\n            >>>\n            >>> # two inputs, two outputs\n            >>> out0, out1 = model(x0, x1)\n            >>>\n            >>> # two inputs, two outputs with return_features=True\n            >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True)\n\n        \"\"\"\n        \n        # forward pass of first input x0\n        f0 = self.backbone(x0).flatten(start_dim=1)\n        z0 = self.projection_mlp(f0)\n        p0 = self.prediction_mlp(z0)\n\n        out0 = (z0, p0)\n\n        # append features if requested\n        if return_features:\n            out0 = (out0, f0)\n\n        # return out0 if x1 is None\n        if x1 is None:\n            return out0\n\n        # forward pass of second input x1\n        f1 = self.backbone(x1).flatten(start_dim=1)\n        z1 = self.projection_mlp(f1)\n        p1 = self.prediction_mlp(z1)\n\n        out1 = (z1, p1)\n\n        # append features if requested\n        if return_features:\n            out1 = (out1, f1)\n\n        # return both outputs\n        return out0, out1\n", "path": "lightly/models/nnclr.py"}]}