prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  CO2 regularizer breaks with PyTorch 1.4.0
  # CO2 regularizer breaks with PyTorch 1.4.0
  
  The `KLDivLoss` in PyTorch only received the `log_target` flag after 1.4.0 so we need to handle the case where the flag is not available.
  
  --- 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/regularizer/co2.py`
  Content:
  ```
  1 """ CO2 Regularizer """
  2 
  3 # Copyright (c) 2020. Lightly AG and its affiliates.
  4 # All Rights Reserved
  5 
  6 import torch
  7 from lightly.loss.memory_bank import MemoryBankModule
  8 
  9 
  10 class CO2Regularizer(MemoryBankModule):
  11     """Implementation of the CO2 regularizer [0] for self-supervised learning.
  12 
  13     [0] CO2, 2021, https://arxiv.org/abs/2010.02217
  14 
  15     Attributes:
  16         alpha:
  17             Weight of the regularization term.
  18         t_consistency:
  19             Temperature used during softmax calculations.
  20         memory_bank_size:
  21             Number of negative samples to store in the memory bank.
  22             Use 0 to use the second batch for negative samples.
  23 
  24     Examples:
  25         >>> # initialize loss function for MoCo
  26         >>> loss_fn = NTXentLoss(memory_bank_size=4096)
  27         >>>
  28         >>> # initialize CO2 regularizer
  29         >>> co2 = CO2Regularizer(alpha=1.0, memory_bank_size=4096)
  30         >>>
  31         >>> # generate two random trasnforms of images
  32         >>> t0 = transforms(images)
  33         >>> t1 = transforms(images)
  34         >>>
  35         >>> # feed through the MoCo model
  36         >>> out0, out1 = model(t0, t1)
  37         >>> 
  38         >>> # calculate loss and apply regularizer
  39         >>> loss = loss_fn(out0, out1) + co2(out0, out1)
  40 
  41     """
  42 
  43     def __init__(self,
  44                 alpha: float = 1,
  45                 t_consistency: float = 0.05,
  46                 memory_bank_size: int = 0):
  47 
  48         super(CO2Regularizer, self).__init__(size=memory_bank_size)
  49         self.kl_div = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)
  50         self.t_consistency = t_consistency
  51         self.alpha = alpha
  52 
  53     def _get_pseudo_labels(self,
  54                            out0: torch.Tensor,
  55                            out1: torch.Tensor,
  56                            negatives: torch.Tensor = None):
  57         """Computes the soft pseudo labels across negative samples.
  58 
  59         Args:
  60             out0:
  61                 Output projections of the first set of transformed images (query).
  62                 Shape: bsz x n_ftrs
  63             out1:
  64                 Output projections of the second set of transformed images (positive sample).
  65                 Shape: bsz x n_ftrs
  66             negatives:
  67                 Negative samples to compare against. If this is None, the second
  68                 batch of images will be used as negative samples.
  69                 Shape: memory_bank_size x n_ftrs
  70 
  71         Returns:
  72             Log probability that a positive samples will classify each negative
  73             sample as the positive sample.
  74             Shape: bsz x (bsz - 1) or bsz x memory_bank_size
  75 
  76         """
  77         batch_size, _ = out0.shape
  78         if negatives is None:
  79             # use second batch as negative samples
  80             # l_pos has shape bsz x 1 and l_neg has shape bsz x bsz
  81             l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)
  82             l_neg = torch.einsum('nc,ck->nk', [out0, out1.t()])
  83             # remove elements on the diagonal
  84             # l_neg has shape bsz x (bsz - 1)
  85             l_neg = l_neg.masked_select(
  86                 ~torch.eye(batch_size, dtype=bool, device=l_neg.device)
  87             ).view(batch_size, batch_size - 1)
  88         else:
  89             # use memory bank as negative samples
  90             # l_pos has shape bsz x 1 and l_neg has shape bsz x memory_bank_size
  91             negatives = negatives.to(out0.device)
  92             l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)
  93             l_neg = torch.einsum('nc,ck->nk', [out0, negatives.clone().detach()])
  94             
  95         # concatenate such that positive samples are at index 0
  96         logits = torch.cat([l_pos, l_neg], dim=1)
  97         # divide by temperature
  98         logits = logits / self.t_consistency
  99 
  100         # the input to kl_div is expected to be log(p) and we set the
  101         # flag log_target to True, so both probabilities should be passed as log
  102         log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
  103         return log_probs
  104 
  105 
  106     def forward(self,
  107                 out0: torch.Tensor,
  108                 out1: torch.Tensor):
  109         """Computes the CO2 regularization term for two model outputs.
  110 
  111         Args:
  112             out0:
  113                 Output projections of the first set of transformed images.
  114             out1:
  115                 Output projections of the second set of transformed images.
  116 
  117         Returns:
  118             The regularization term multiplied by the weight factor alpha.
  119 
  120         """
  121 
  122         # normalize the output to length 1
  123         out0 = torch.nn.functional.normalize(out0, dim=1)
  124         out1 = torch.nn.functional.normalize(out1, dim=1)
  125 
  126         # ask memory bank for negative samples and extend it with out1 if 
  127         # out1 requires a gradient, otherwise keep the same vectors in the 
  128         # memory bank (this allows for keeping the memory bank constant e.g.
  129         # for evaluating the loss on the test set)
  130         # if the memory_bank size is 0, negatives will be None
  131         out1, negatives = \
  132             super(CO2Regularizer, self).forward(out1, update=True)
  133         
  134         # get log probabilities
  135         p = self._get_pseudo_labels(out0, out1, negatives)
  136         q = self._get_pseudo_labels(out1, out0, negatives)
  137         
  138         # calculate kullback leibler divergence from log probabilities
  139         return self.alpha * 0.5 * (self.kl_div(p, q) + self.kl_div(q, p))
  140
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/lightly/loss/regularizer/co2.py b/lightly/loss/regularizer/co2.py
  --- a/lightly/loss/regularizer/co2.py
  +++ b/lightly/loss/regularizer/co2.py
  @@ -46,7 +46,19 @@
                   memory_bank_size: int = 0):
   
           super(CO2Regularizer, self).__init__(size=memory_bank_size)
  -        self.kl_div = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)
  +        # try-catch the KLDivLoss construction for backwards compatability
  +        self.log_target = True
  +        try:
  +            self.kl_div = torch.nn.KLDivLoss(
  +                reduction='batchmean',
  +                log_target=True
  +            )
  +        except TypeError:
  +            self.log_target = False
  +            self.kl_div = torch.nn.KLDivLoss(
  +                reduction='batchmean'
  +            )
  +
           self.t_consistency = t_consistency
           self.alpha = alpha
   
  @@ -97,10 +109,8 @@
           # divide by temperature
           logits = logits / self.t_consistency
   
  -        # the input to kl_div is expected to be log(p) and we set the
  -        # flag log_target to True, so both probabilities should be passed as log
  -        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
  -        return log_probs
  +        # the input to kl_div is expected to be log(p) 
  +        return torch.nn.functional.log_softmax(logits, dim=-1)
   
   
       def forward(self,
  @@ -135,5 +145,11 @@
           p = self._get_pseudo_labels(out0, out1, negatives)
           q = self._get_pseudo_labels(out1, out0, negatives)
           
  -        # calculate kullback leibler divergence from log probabilities
  -        return self.alpha * 0.5 * (self.kl_div(p, q) + self.kl_div(q, p))
  +        # calculate symmetrized kullback leibler divergence
  +        if self.log_target:
  +            div = self.kl_div(p, q) + self.kl_div(q, p)
  +        else:
  +            # can't use log_target because of early torch version
  +            div = self.kl_div(p, torch.exp(q)) + self.kl_div(q, torch.exp(p))
  +
  +        return self.alpha * 0.5 * div
  
verification_info: {"golden_diff": "diff --git a/lightly/loss/regularizer/co2.py b/lightly/loss/regularizer/co2.py\n--- a/lightly/loss/regularizer/co2.py\n+++ b/lightly/loss/regularizer/co2.py\n@@ -46,7 +46,19 @@\n                 memory_bank_size: int = 0):\n \n         super(CO2Regularizer, self).__init__(size=memory_bank_size)\n-        self.kl_div = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)\n+        # try-catch the KLDivLoss construction for backwards compatability\n+        self.log_target = True\n+        try:\n+            self.kl_div = torch.nn.KLDivLoss(\n+                reduction='batchmean',\n+                log_target=True\n+            )\n+        except TypeError:\n+            self.log_target = False\n+            self.kl_div = torch.nn.KLDivLoss(\n+                reduction='batchmean'\n+            )\n+\n         self.t_consistency = t_consistency\n         self.alpha = alpha\n \n@@ -97,10 +109,8 @@\n         # divide by temperature\n         logits = logits / self.t_consistency\n \n-        # the input to kl_div is expected to be log(p) and we set the\n-        # flag log_target to True, so both probabilities should be passed as log\n-        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)\n-        return log_probs\n+        # the input to kl_div is expected to be log(p) \n+        return torch.nn.functional.log_softmax(logits, dim=-1)\n \n \n     def forward(self,\n@@ -135,5 +145,11 @@\n         p = self._get_pseudo_labels(out0, out1, negatives)\n         q = self._get_pseudo_labels(out1, out0, negatives)\n         \n-        # calculate kullback leibler divergence from log probabilities\n-        return self.alpha * 0.5 * (self.kl_div(p, q) + self.kl_div(q, p))\n+        # calculate symmetrized kullback leibler divergence\n+        if self.log_target:\n+            div = self.kl_div(p, q) + self.kl_div(q, p)\n+        else:\n+            # can't use log_target because of early torch version\n+            div = self.kl_div(p, torch.exp(q)) + self.kl_div(q, torch.exp(p))\n+\n+        return self.alpha * 0.5 * div\n", "issue": "CO2 regularizer breaks with PyTorch 1.4.0\n# CO2 regularizer breaks with PyTorch 1.4.0\r\n\r\nThe `KLDivLoss` in PyTorch only received the `log_target` flag after 1.4.0 so we need to handle the case where the flag is not available.\n", "before_files": [{"content": "\"\"\" CO2 Regularizer \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nfrom lightly.loss.memory_bank import MemoryBankModule\n\n\nclass CO2Regularizer(MemoryBankModule):\n    \"\"\"Implementation of the CO2 regularizer [0] for self-supervised learning.\n\n    [0] CO2, 2021, https://arxiv.org/abs/2010.02217\n\n    Attributes:\n        alpha:\n            Weight of the regularization term.\n        t_consistency:\n            Temperature used during softmax calculations.\n        memory_bank_size:\n            Number of negative samples to store in the memory bank.\n            Use 0 to use the second batch for negative samples.\n\n    Examples:\n        >>> # initialize loss function for MoCo\n        >>> loss_fn = NTXentLoss(memory_bank_size=4096)\n        >>>\n        >>> # initialize CO2 regularizer\n        >>> co2 = CO2Regularizer(alpha=1.0, memory_bank_size=4096)\n        >>>\n        >>> # generate two random trasnforms of images\n        >>> t0 = transforms(images)\n        >>> t1 = transforms(images)\n        >>>\n        >>> # feed through the MoCo model\n        >>> out0, out1 = model(t0, t1)\n        >>> \n        >>> # calculate loss and apply regularizer\n        >>> loss = loss_fn(out0, out1) + co2(out0, out1)\n\n    \"\"\"\n\n    def __init__(self,\n                alpha: float = 1,\n                t_consistency: float = 0.05,\n                memory_bank_size: int = 0):\n\n        super(CO2Regularizer, self).__init__(size=memory_bank_size)\n        self.kl_div = torch.nn.KLDivLoss(reduction='batchmean', log_target=True)\n        self.t_consistency = t_consistency\n        self.alpha = alpha\n\n    def _get_pseudo_labels(self,\n                           out0: torch.Tensor,\n                           out1: torch.Tensor,\n                           negatives: torch.Tensor = None):\n        \"\"\"Computes the soft pseudo labels across negative samples.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images (query).\n                Shape: bsz x n_ftrs\n            out1:\n                Output projections of the second set of transformed images (positive sample).\n                Shape: bsz x n_ftrs\n            negatives:\n                Negative samples to compare against. If this is None, the second\n                batch of images will be used as negative samples.\n                Shape: memory_bank_size x n_ftrs\n\n        Returns:\n            Log probability that a positive samples will classify each negative\n            sample as the positive sample.\n            Shape: bsz x (bsz - 1) or bsz x memory_bank_size\n\n        \"\"\"\n        batch_size, _ = out0.shape\n        if negatives is None:\n            # use second batch as negative samples\n            # l_pos has shape bsz x 1 and l_neg has shape bsz x bsz\n            l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)\n            l_neg = torch.einsum('nc,ck->nk', [out0, out1.t()])\n            # remove elements on the diagonal\n            # l_neg has shape bsz x (bsz - 1)\n            l_neg = l_neg.masked_select(\n                ~torch.eye(batch_size, dtype=bool, device=l_neg.device)\n            ).view(batch_size, batch_size - 1)\n        else:\n            # use memory bank as negative samples\n            # l_pos has shape bsz x 1 and l_neg has shape bsz x memory_bank_size\n            negatives = negatives.to(out0.device)\n            l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)\n            l_neg = torch.einsum('nc,ck->nk', [out0, negatives.clone().detach()])\n            \n        # concatenate such that positive samples are at index 0\n        logits = torch.cat([l_pos, l_neg], dim=1)\n        # divide by temperature\n        logits = logits / self.t_consistency\n\n        # the input to kl_div is expected to be log(p) and we set the\n        # flag log_target to True, so both probabilities should be passed as log\n        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)\n        return log_probs\n\n\n    def forward(self,\n                out0: torch.Tensor,\n                out1: torch.Tensor):\n        \"\"\"Computes the CO2 regularization term for two model outputs.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images.\n            out1:\n                Output projections of the second set of transformed images.\n\n        Returns:\n            The regularization term multiplied by the weight factor alpha.\n\n        \"\"\"\n\n        # normalize the output to length 1\n        out0 = torch.nn.functional.normalize(out0, dim=1)\n        out1 = torch.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        # if the memory_bank size is 0, negatives will be None\n        out1, negatives = \\\n            super(CO2Regularizer, self).forward(out1, update=True)\n        \n        # get log probabilities\n        p = self._get_pseudo_labels(out0, out1, negatives)\n        q = self._get_pseudo_labels(out1, out0, negatives)\n        \n        # calculate kullback leibler divergence from log probabilities\n        return self.alpha * 0.5 * (self.kl_div(p, q) + self.kl_div(q, p))\n", "path": "lightly/loss/regularizer/co2.py"}], "after_files": [{"content": "\"\"\" CO2 Regularizer \"\"\"\n\n# Copyright (c) 2020. Lightly AG and its affiliates.\n# All Rights Reserved\n\nimport torch\nfrom lightly.loss.memory_bank import MemoryBankModule\n\n\nclass CO2Regularizer(MemoryBankModule):\n    \"\"\"Implementation of the CO2 regularizer [0] for self-supervised learning.\n\n    [0] CO2, 2021, https://arxiv.org/abs/2010.02217\n\n    Attributes:\n        alpha:\n            Weight of the regularization term.\n        t_consistency:\n            Temperature used during softmax calculations.\n        memory_bank_size:\n            Number of negative samples to store in the memory bank.\n            Use 0 to use the second batch for negative samples.\n\n    Examples:\n        >>> # initialize loss function for MoCo\n        >>> loss_fn = NTXentLoss(memory_bank_size=4096)\n        >>>\n        >>> # initialize CO2 regularizer\n        >>> co2 = CO2Regularizer(alpha=1.0, memory_bank_size=4096)\n        >>>\n        >>> # generate two random trasnforms of images\n        >>> t0 = transforms(images)\n        >>> t1 = transforms(images)\n        >>>\n        >>> # feed through the MoCo model\n        >>> out0, out1 = model(t0, t1)\n        >>> \n        >>> # calculate loss and apply regularizer\n        >>> loss = loss_fn(out0, out1) + co2(out0, out1)\n\n    \"\"\"\n\n    def __init__(self,\n                alpha: float = 1,\n                t_consistency: float = 0.05,\n                memory_bank_size: int = 0):\n\n        super(CO2Regularizer, self).__init__(size=memory_bank_size)\n        # try-catch the KLDivLoss construction for backwards compatability\n        self.log_target = True\n        try:\n            self.kl_div = torch.nn.KLDivLoss(\n                reduction='batchmean',\n                log_target=True\n            )\n        except TypeError:\n            self.log_target = False\n            self.kl_div = torch.nn.KLDivLoss(\n                reduction='batchmean'\n            )\n\n        self.t_consistency = t_consistency\n        self.alpha = alpha\n\n    def _get_pseudo_labels(self,\n                           out0: torch.Tensor,\n                           out1: torch.Tensor,\n                           negatives: torch.Tensor = None):\n        \"\"\"Computes the soft pseudo labels across negative samples.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images (query).\n                Shape: bsz x n_ftrs\n            out1:\n                Output projections of the second set of transformed images (positive sample).\n                Shape: bsz x n_ftrs\n            negatives:\n                Negative samples to compare against. If this is None, the second\n                batch of images will be used as negative samples.\n                Shape: memory_bank_size x n_ftrs\n\n        Returns:\n            Log probability that a positive samples will classify each negative\n            sample as the positive sample.\n            Shape: bsz x (bsz - 1) or bsz x memory_bank_size\n\n        \"\"\"\n        batch_size, _ = out0.shape\n        if negatives is None:\n            # use second batch as negative samples\n            # l_pos has shape bsz x 1 and l_neg has shape bsz x bsz\n            l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)\n            l_neg = torch.einsum('nc,ck->nk', [out0, out1.t()])\n            # remove elements on the diagonal\n            # l_neg has shape bsz x (bsz - 1)\n            l_neg = l_neg.masked_select(\n                ~torch.eye(batch_size, dtype=bool, device=l_neg.device)\n            ).view(batch_size, batch_size - 1)\n        else:\n            # use memory bank as negative samples\n            # l_pos has shape bsz x 1 and l_neg has shape bsz x memory_bank_size\n            negatives = negatives.to(out0.device)\n            l_pos = torch.einsum('nc,nc->n', [out0, out1]).unsqueeze(-1)\n            l_neg = torch.einsum('nc,ck->nk', [out0, negatives.clone().detach()])\n            \n        # concatenate such that positive samples are at index 0\n        logits = torch.cat([l_pos, l_neg], dim=1)\n        # divide by temperature\n        logits = logits / self.t_consistency\n\n        # the input to kl_div is expected to be log(p) \n        return torch.nn.functional.log_softmax(logits, dim=-1)\n\n\n    def forward(self,\n                out0: torch.Tensor,\n                out1: torch.Tensor):\n        \"\"\"Computes the CO2 regularization term for two model outputs.\n\n        Args:\n            out0:\n                Output projections of the first set of transformed images.\n            out1:\n                Output projections of the second set of transformed images.\n\n        Returns:\n            The regularization term multiplied by the weight factor alpha.\n\n        \"\"\"\n\n        # normalize the output to length 1\n        out0 = torch.nn.functional.normalize(out0, dim=1)\n        out1 = torch.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        # if the memory_bank size is 0, negatives will be None\n        out1, negatives = \\\n            super(CO2Regularizer, self).forward(out1, update=True)\n        \n        # get log probabilities\n        p = self._get_pseudo_labels(out0, out1, negatives)\n        q = self._get_pseudo_labels(out1, out0, negatives)\n        \n        # calculate symmetrized kullback leibler divergence\n        if self.log_target:\n            div = self.kl_div(p, q) + self.kl_div(q, p)\n        else:\n            # can't use log_target because of early torch version\n            div = self.kl_div(p, torch.exp(q)) + self.kl_div(q, torch.exp(p))\n\n        return self.alpha * 0.5 * div\n", "path": "lightly/loss/regularizer/co2.py"}]}