We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Bug in PMSNLoss
Nice implementation of the [PMSNLoss](https://github.com/lightly-ai/lightly/blob/ddfed3c4dc03a8d2722df24bfa537d24ac80bde6/lightly/loss/pmsn_loss.py)! But the computation of Kullback-Leibler divergence missed `.log()` in Line 71&142.


--- 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/pmsn_loss.py`
Content:
```
1 from typing import Callable
2 
3 import torch
4 import torch.nn.functional as F
5 from torch import Tensor
6 
7 from lightly.loss.msn_loss import MSNLoss
8 
9 
10 class PMSNLoss(MSNLoss):
11     """Implementation of the loss function from PMSN [0] using a power law target
12     distribution.
13 
14     - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277
15 
16     Attributes:
17         temperature:
18             Similarities between anchors and targets are scaled by the inverse of
19             the temperature. Must be in (0, inf).
20         sinkhorn_iterations:
21             Number of sinkhorn normalization iterations on the targets.
22         regularization_weight:
23             Weight factor lambda by which the regularization loss is scaled. Set to 0
24             to disable regularization.
25         power_law_exponent:
26             Exponent for power law distribution. Entry k of the distribution is
27             proportional to (1 / k) ^ power_law_exponent, with k ranging from 1 to dim + 1.
28         gather_distributed:
29             If True, then target probabilities are gathered from all GPUs.
30 
31      Examples:
32 
33         >>> # initialize loss function
34         >>> loss_fn = PMSNLoss()
35         >>>
36         >>> # generate anchors and targets of images
37         >>> anchors = transforms(images)
38         >>> targets = transforms(images)
39         >>>
40         >>> # feed through PMSN model
41         >>> anchors_out = model(anchors)
42         >>> targets_out = model.target(targets)
43         >>>
44         >>> # calculate loss
45         >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)
46     """
47 
48     def __init__(
49         self,
50         temperature: float = 0.1,
51         sinkhorn_iterations: int = 3,
52         regularization_weight: float = 1,
53         power_law_exponent: float = 0.25,
54         gather_distributed: bool = False,
55     ):
56         super().__init__(
57             temperature=temperature,
58             sinkhorn_iterations=sinkhorn_iterations,
59             regularization_weight=regularization_weight,
60             gather_distributed=gather_distributed,
61         )
62         self.power_law_exponent = power_law_exponent
63 
64     def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:
65         """Calculates regularization loss with a power law target distribution."""
66         power_dist = _power_law_distribution(
67             size=mean_anchor_probs.shape[0],
68             exponent=self.power_law_exponent,
69             device=mean_anchor_probs.device,
70         )
71         loss = F.kl_div(input=mean_anchor_probs, target=power_dist, reduction="sum")
72         return loss
73 
74 
75 class PMSNCustomLoss(MSNLoss):
76     """Implementation of the loss function from PMSN [0] with a custom target
77     distribution.
78 
79     - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277
80 
81     Attributes:
82         target_distribution:
83             A function that takes the mean anchor probabilities tensor with shape (dim,)
84             as input and returns a target probability distribution tensor with the same
85             shape. The returned distribution should sum up to one. The final
86             regularization loss is calculated as KL(mean_anchor_probs, target_dist)
87             where KL is the Kullback-Leibler divergence.
88         temperature:
89             Similarities between anchors and targets are scaled by the inverse of
90             the temperature. Must be in (0, inf).
91         sinkhorn_iterations:
92             Number of sinkhorn normalization iterations on the targets.
93         regularization_weight:
94             Weight factor lambda by which the regularization loss is scaled. Set to 0
95             to disable regularization.
96         gather_distributed:
97             If True, then target probabilities are gathered from all GPUs.
98 
99      Examples:
100 
101         >>> # define custom target distribution
102         >>> def my_uniform_distribution(mean_anchor_probabilities: Tensor) -> Tensor:
103         >>>     dim = mean_anchor_probabilities.shape[0]
104         >>>     return mean_anchor_probabilities.new_ones(dim) / dim
105         >>>
106         >>> # initialize loss function
107         >>> loss_fn = PMSNCustomLoss(target_distribution=my_uniform_distribution)
108         >>>
109         >>> # generate anchors and targets of images
110         >>> anchors = transforms(images)
111         >>> targets = transforms(images)
112         >>>
113         >>> # feed through PMSN model
114         >>> anchors_out = model(anchors)
115         >>> targets_out = model.target(targets)
116         >>>
117         >>> # calculate loss
118         >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)
119     """
120 
121     def __init__(
122         self,
123         target_distribution: Callable[[Tensor], Tensor],
124         temperature: float = 0.1,
125         sinkhorn_iterations: int = 3,
126         regularization_weight: float = 1,
127         gather_distributed: bool = False,
128     ):
129         super().__init__(
130             temperature=temperature,
131             sinkhorn_iterations=sinkhorn_iterations,
132             regularization_weight=regularization_weight,
133             gather_distributed=gather_distributed,
134         )
135         self.target_distribution = target_distribution
136 
137     def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:
138         """Calculates regularization loss with a custom target distribution."""
139         target_dist = self.target_distribution(mean_anchor_probs).to(
140             mean_anchor_probs.device
141         )
142         loss = F.kl_div(input=mean_anchor_probs, target=target_dist, reduction="sum")
143         return loss
144 
145 
146 def _power_law_distribution(size: int, exponent: float, device: torch.device) -> Tensor:
147     """Returns a power law distribution summing up to 1."""
148     k = torch.arange(1, size + 1, device=device)
149     power_dist = k ** (-exponent)
150     power_dist = power_dist / power_dist.sum()
151     return power_dist
152
```
--- 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/pmsn_loss.py b/lightly/loss/pmsn_loss.py
--- a/lightly/loss/pmsn_loss.py
+++ b/lightly/loss/pmsn_loss.py
@@ -68,7 +68,9 @@
             exponent=self.power_law_exponent,
             device=mean_anchor_probs.device,
         )
-        loss = F.kl_div(input=mean_anchor_probs, target=power_dist, reduction="sum")
+        loss = F.kl_div(
+            input=mean_anchor_probs.log(), target=power_dist, reduction="sum"
+        )
         return loss
 
 
@@ -139,7 +141,9 @@
         target_dist = self.target_distribution(mean_anchor_probs).to(
             mean_anchor_probs.device
         )
-        loss = F.kl_div(input=mean_anchor_probs, target=target_dist, reduction="sum")
+        loss = F.kl_div(
+            input=mean_anchor_probs.log(), target=target_dist, reduction="sum"
+        )
         return loss

{"golden_diff": "diff --git a/lightly/loss/pmsn_loss.py b/lightly/loss/pmsn_loss.py\n--- a/lightly/loss/pmsn_loss.py\n+++ b/lightly/loss/pmsn_loss.py\n@@ -68,7 +68,9 @@\n             exponent=self.power_law_exponent,\n             device=mean_anchor_probs.device,\n         )\n-        loss = F.kl_div(input=mean_anchor_probs, target=power_dist, reduction=\"sum\")\n+        loss = F.kl_div(\n+            input=mean_anchor_probs.log(), target=power_dist, reduction=\"sum\"\n+        )\n         return loss\n \n \n@@ -139,7 +141,9 @@\n         target_dist = self.target_distribution(mean_anchor_probs).to(\n             mean_anchor_probs.device\n         )\n-        loss = F.kl_div(input=mean_anchor_probs, target=target_dist, reduction=\"sum\")\n+        loss = F.kl_div(\n+            input=mean_anchor_probs.log(), target=target_dist, reduction=\"sum\"\n+        )\n         return loss\n", "issue": "Bug in PMSNLoss\nNice implementation of the [PMSNLoss](https://github.com/lightly-ai/lightly/blob/ddfed3c4dc03a8d2722df24bfa537d24ac80bde6/lightly/loss/pmsn_loss.py)! But the computation of Kullback-Leibler divergence missed `.log()` in Line 71&142.\r\n\n", "before_files": [{"content": "from typing import Callable\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import Tensor\n\nfrom lightly.loss.msn_loss import MSNLoss\n\n\nclass PMSNLoss(MSNLoss):\n    \"\"\"Implementation of the loss function from PMSN [0] using a power law target\n    distribution.\n\n    - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277\n\n    Attributes:\n        temperature:\n            Similarities between anchors and targets are scaled by the inverse of\n            the temperature. Must be in (0, inf).\n        sinkhorn_iterations:\n            Number of sinkhorn normalization iterations on the targets.\n        regularization_weight:\n            Weight factor lambda by which the regularization loss is scaled. Set to 0\n            to disable regularization.\n        power_law_exponent:\n            Exponent for power law distribution. Entry k of the distribution is\n            proportional to (1 / k) ^ power_law_exponent, with k ranging from 1 to dim + 1.\n        gather_distributed:\n            If True, then target probabilities are gathered from all GPUs.\n\n     Examples:\n\n        >>> # initialize loss function\n        >>> loss_fn = PMSNLoss()\n        >>>\n        >>> # generate anchors and targets of images\n        >>> anchors = transforms(images)\n        >>> targets = transforms(images)\n        >>>\n        >>> # feed through PMSN model\n        >>> anchors_out = model(anchors)\n        >>> targets_out = model.target(targets)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)\n    \"\"\"\n\n    def __init__(\n        self,\n        temperature: float = 0.1,\n        sinkhorn_iterations: int = 3,\n        regularization_weight: float = 1,\n        power_law_exponent: float = 0.25,\n        gather_distributed: bool = False,\n    ):\n        super().__init__(\n            temperature=temperature,\n            sinkhorn_iterations=sinkhorn_iterations,\n            regularization_weight=regularization_weight,\n            gather_distributed=gather_distributed,\n        )\n        self.power_law_exponent = power_law_exponent\n\n    def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:\n        \"\"\"Calculates regularization loss with a power law target distribution.\"\"\"\n        power_dist = _power_law_distribution(\n            size=mean_anchor_probs.shape[0],\n            exponent=self.power_law_exponent,\n            device=mean_anchor_probs.device,\n        )\n        loss = F.kl_div(input=mean_anchor_probs, target=power_dist, reduction=\"sum\")\n        return loss\n\n\nclass PMSNCustomLoss(MSNLoss):\n    \"\"\"Implementation of the loss function from PMSN [0] with a custom target\n    distribution.\n\n    - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277\n\n    Attributes:\n        target_distribution:\n            A function that takes the mean anchor probabilities tensor with shape (dim,)\n            as input and returns a target probability distribution tensor with the same\n            shape. The returned distribution should sum up to one. The final\n            regularization loss is calculated as KL(mean_anchor_probs, target_dist)\n            where KL is the Kullback-Leibler divergence.\n        temperature:\n            Similarities between anchors and targets are scaled by the inverse of\n            the temperature. Must be in (0, inf).\n        sinkhorn_iterations:\n            Number of sinkhorn normalization iterations on the targets.\n        regularization_weight:\n            Weight factor lambda by which the regularization loss is scaled. Set to 0\n            to disable regularization.\n        gather_distributed:\n            If True, then target probabilities are gathered from all GPUs.\n\n     Examples:\n\n        >>> # define custom target distribution\n        >>> def my_uniform_distribution(mean_anchor_probabilities: Tensor) -> Tensor:\n        >>>     dim = mean_anchor_probabilities.shape[0]\n        >>>     return mean_anchor_probabilities.new_ones(dim) / dim\n        >>>\n        >>> # initialize loss function\n        >>> loss_fn = PMSNCustomLoss(target_distribution=my_uniform_distribution)\n        >>>\n        >>> # generate anchors and targets of images\n        >>> anchors = transforms(images)\n        >>> targets = transforms(images)\n        >>>\n        >>> # feed through PMSN model\n        >>> anchors_out = model(anchors)\n        >>> targets_out = model.target(targets)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)\n    \"\"\"\n\n    def __init__(\n        self,\n        target_distribution: Callable[[Tensor], Tensor],\n        temperature: float = 0.1,\n        sinkhorn_iterations: int = 3,\n        regularization_weight: float = 1,\n        gather_distributed: bool = False,\n    ):\n        super().__init__(\n            temperature=temperature,\n            sinkhorn_iterations=sinkhorn_iterations,\n            regularization_weight=regularization_weight,\n            gather_distributed=gather_distributed,\n        )\n        self.target_distribution = target_distribution\n\n    def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:\n        \"\"\"Calculates regularization loss with a custom target distribution.\"\"\"\n        target_dist = self.target_distribution(mean_anchor_probs).to(\n            mean_anchor_probs.device\n        )\n        loss = F.kl_div(input=mean_anchor_probs, target=target_dist, reduction=\"sum\")\n        return loss\n\n\ndef _power_law_distribution(size: int, exponent: float, device: torch.device) -> Tensor:\n    \"\"\"Returns a power law distribution summing up to 1.\"\"\"\n    k = torch.arange(1, size + 1, device=device)\n    power_dist = k ** (-exponent)\n    power_dist = power_dist / power_dist.sum()\n    return power_dist\n", "path": "lightly/loss/pmsn_loss.py"}], "after_files": [{"content": "from typing import Callable\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import Tensor\n\nfrom lightly.loss.msn_loss import MSNLoss\n\n\nclass PMSNLoss(MSNLoss):\n    \"\"\"Implementation of the loss function from PMSN [0] using a power law target\n    distribution.\n\n    - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277\n\n    Attributes:\n        temperature:\n            Similarities between anchors and targets are scaled by the inverse of\n            the temperature. Must be in (0, inf).\n        sinkhorn_iterations:\n            Number of sinkhorn normalization iterations on the targets.\n        regularization_weight:\n            Weight factor lambda by which the regularization loss is scaled. Set to 0\n            to disable regularization.\n        power_law_exponent:\n            Exponent for power law distribution. Entry k of the distribution is\n            proportional to (1 / k) ^ power_law_exponent, with k ranging from 1 to dim + 1.\n        gather_distributed:\n            If True, then target probabilities are gathered from all GPUs.\n\n     Examples:\n\n        >>> # initialize loss function\n        >>> loss_fn = PMSNLoss()\n        >>>\n        >>> # generate anchors and targets of images\n        >>> anchors = transforms(images)\n        >>> targets = transforms(images)\n        >>>\n        >>> # feed through PMSN model\n        >>> anchors_out = model(anchors)\n        >>> targets_out = model.target(targets)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)\n    \"\"\"\n\n    def __init__(\n        self,\n        temperature: float = 0.1,\n        sinkhorn_iterations: int = 3,\n        regularization_weight: float = 1,\n        power_law_exponent: float = 0.25,\n        gather_distributed: bool = False,\n    ):\n        super().__init__(\n            temperature=temperature,\n            sinkhorn_iterations=sinkhorn_iterations,\n            regularization_weight=regularization_weight,\n            gather_distributed=gather_distributed,\n        )\n        self.power_law_exponent = power_law_exponent\n\n    def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:\n        \"\"\"Calculates regularization loss with a power law target distribution.\"\"\"\n        power_dist = _power_law_distribution(\n            size=mean_anchor_probs.shape[0],\n            exponent=self.power_law_exponent,\n            device=mean_anchor_probs.device,\n        )\n        loss = F.kl_div(\n            input=mean_anchor_probs.log(), target=power_dist, reduction=\"sum\"\n        )\n        return loss\n\n\nclass PMSNCustomLoss(MSNLoss):\n    \"\"\"Implementation of the loss function from PMSN [0] with a custom target\n    distribution.\n\n    - [0]: Prior Matching for Siamese Networks, 2022, https://arxiv.org/abs/2210.07277\n\n    Attributes:\n        target_distribution:\n            A function that takes the mean anchor probabilities tensor with shape (dim,)\n            as input and returns a target probability distribution tensor with the same\n            shape. The returned distribution should sum up to one. The final\n            regularization loss is calculated as KL(mean_anchor_probs, target_dist)\n            where KL is the Kullback-Leibler divergence.\n        temperature:\n            Similarities between anchors and targets are scaled by the inverse of\n            the temperature. Must be in (0, inf).\n        sinkhorn_iterations:\n            Number of sinkhorn normalization iterations on the targets.\n        regularization_weight:\n            Weight factor lambda by which the regularization loss is scaled. Set to 0\n            to disable regularization.\n        gather_distributed:\n            If True, then target probabilities are gathered from all GPUs.\n\n     Examples:\n\n        >>> # define custom target distribution\n        >>> def my_uniform_distribution(mean_anchor_probabilities: Tensor) -> Tensor:\n        >>>     dim = mean_anchor_probabilities.shape[0]\n        >>>     return mean_anchor_probabilities.new_ones(dim) / dim\n        >>>\n        >>> # initialize loss function\n        >>> loss_fn = PMSNCustomLoss(target_distribution=my_uniform_distribution)\n        >>>\n        >>> # generate anchors and targets of images\n        >>> anchors = transforms(images)\n        >>> targets = transforms(images)\n        >>>\n        >>> # feed through PMSN model\n        >>> anchors_out = model(anchors)\n        >>> targets_out = model.target(targets)\n        >>>\n        >>> # calculate loss\n        >>> loss = loss_fn(anchors_out, targets_out, prototypes=model.prototypes)\n    \"\"\"\n\n    def __init__(\n        self,\n        target_distribution: Callable[[Tensor], Tensor],\n        temperature: float = 0.1,\n        sinkhorn_iterations: int = 3,\n        regularization_weight: float = 1,\n        gather_distributed: bool = False,\n    ):\n        super().__init__(\n            temperature=temperature,\n            sinkhorn_iterations=sinkhorn_iterations,\n            regularization_weight=regularization_weight,\n            gather_distributed=gather_distributed,\n        )\n        self.target_distribution = target_distribution\n\n    def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:\n        \"\"\"Calculates regularization loss with a custom target distribution.\"\"\"\n        target_dist = self.target_distribution(mean_anchor_probs).to(\n            mean_anchor_probs.device\n        )\n        loss = F.kl_div(\n            input=mean_anchor_probs.log(), target=target_dist, reduction=\"sum\"\n        )\n        return loss\n\n\ndef _power_law_distribution(size: int, exponent: float, device: torch.device) -> Tensor:\n    \"\"\"Returns a power law distribution summing up to 1.\"\"\"\n    k = torch.arange(1, size + 1, device=device)\n    power_dist = k ** (-exponent)\n    power_dist = power_dist / power_dist.sum()\n    return power_dist\n", "path": "lightly/loss/pmsn_loss.py"}]}