internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: lightly-ai/lightly

Issue:
LightlyDataset for videos has squared complexity
## Description
When creating/using a video dataset, it calls the function `VideoDataset.get_filename(self, index)`/`VideoDataset__getitem__(self, index)` for every frame. Both of these methods do an index check and thus call `self.__len__()`. 
https://github.com/lightly-ai/lightly/blob/592d8f899a427dd0763b66ef856c8ffd74179357/lightly/data/_video.py#L334-L338

However, this method iterates over all `video_timestamps`. Thus, the index check has a squared complexity w.r.t the number of videos. 

I could reproduce this squared complexity with measurements on datasets and with profiling.


## Task
fix the squared complexity


Discussion:

Resolving patch:
diff --git a/lightly/data/_video.py b/lightly/data/_video.py
--- a/lightly/data/_video.py
+++ b/lightly/data/_video.py
@@ -271,6 +271,9 @@ def __init__(self,
 
         self.videos = videos
         self.video_timestamps = video_timestamps
+        self._length = sum((
+            len(ts) for ts in self.video_timestamps
+        ))
         # offsets[i] indicates the index of the first frame of the i-th video.
         # e.g. for two videos of length 10 and 20, the offsets will be [0, 10].
         self.offsets = offsets
@@ -334,8 +337,10 @@ def __getitem__(self, index):
     def __len__(self):
         """Returns the number of samples (frames) in the dataset.
 
+        This can be precomputed, because self.video_timestamps is only
+        set in the __init__
         """
-        return sum((len(ts) for ts in self.video_timestamps))
+        return self._length
 
     def get_filename(self, index):
         """Returns a filename for the frame at index.


Test patch:
Repository: lightly-ai/lightly

Issue:
Rewrite Docs with Low-Level Building Blocks: Python API
# Rewrite Docs with Low-Level Building Blocks: Main Concepts

The [Python API](https://docs.lightly.ai/lightly.html) pages still contain mid- and high-level building blocks.

- [x] Remove the documentation for the mid-level building blocks
- [x] Remove the documentation for the high-level building blocks

Related to #493.


Discussion:

Resolving patch:
diff --git a/lightly/__init__.py b/lightly/__init__.py
--- a/lightly/__init__.py
+++ b/lightly/__init__.py
@@ -29,11 +29,6 @@
   collate functions are in charge of the data augmentations which are crucial for
   self-supervised learning.
 
-- **embedding**:
-
-  The lightly.embedding module combines the self-supervised models with a dataloader,
-  optimizer, and loss function to provide a simple pytorch-lightning trainable.
-
 - **loss**:
 
   The lightly.loss module contains implementations of popular self-supervised training
@@ -41,20 +36,22 @@
 
 - **models**:
 
-  The lightly.models module holds the implementation of the ResNet as well as self-
-  supervised methods. Currently implements:
-
-  - SimCLR
-
-  - MoCo
-
-  - SimSiam
+  The lightly.models module holds the implementation of the ResNet as well as heads
+  for self-supervised methods. It currently implements the heads of:
 
   - Barlow Twins
 
   - BYOL
-
+  
+  - MoCo
+  
   - NNCLR
+  
+  - SimCLR
+  
+  - SimSiam
+  
+  - SwaV
 
 - **transforms**:
 


Test patch:
Repository: lightly-ai/lightly

Issue:
Update SimSiam module in SimSiam tutorial
The SimSiam module in the SimSiam tutorial requires currently `x0` and `x1` in the forward pass but they are independent. We can simplify the forward function to accept only a single `x`.


Discussion:

Resolving patch:
diff --git a/docs/source/tutorials_source/package/tutorial_simsiam_esa.py b/docs/source/tutorials_source/package/tutorial_simsiam_esa.py
--- a/docs/source/tutorials_source/package/tutorial_simsiam_esa.py
+++ b/docs/source/tutorials_source/package/tutorial_simsiam_esa.py
@@ -167,20 +167,16 @@ def __init__(
             out_dim, pred_hidden_dim, out_dim
         )
 
-    def forward(self, x0, x1):
+    def forward(self, x):
         # get representations
-        f0 = self.backbone(x0).flatten(start_dim=1)
-        f1 = self.backbone(x1).flatten(start_dim=1)
+        f = self.backbone(x).flatten(start_dim=1)
         # get projections
-        z0 = self.projection_head(f0)
-        z1 = self.projection_head(f1)
+        z = self.projection_head(f)
         # get predictions
-        p0 = self.prediction_head(z0)
-        p1 = self.prediction_head(z1)
+        p = self.prediction_head(z)
         # stop gradient
-        z0 = z0.detach()
-        z1 = z1.detach()
-        return p0, p1, z0, z1
+        z = z.detach()
+        return z, p
 
 
 # we use a pretrained resnet for this tutorial to speed
@@ -238,13 +234,14 @@ def forward(self, x0, x1):
         x1 = x1.to(device)
 
         # run the model on both transforms of the images
-        # we get predictions (p0 and p1) and 
-        # projections (z0 and z1) as output
-        p0, p1, z0, z1 = model(x0, x1)
+        # we get projections (z0 and z1) and
+        # predictions (p0 and p1) as output
+        z0, p0 = model(x0)
+        z1, p1 = model(x1)
 
         # apply the symmetric negative cosine similarity
         # and run backpropagation
-        loss = 0.5 * (criterion(p0, z0) + criterion(p1, z1))
+        loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))
         loss.backward()
 
         optimizer.step()


Test patch:
Repository: lightly-ai/lightly

Issue:
Update docs version
Update docs version and link copyright in docs footer to lightly website

Closes #618 


Discussion:

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -15,16 +15,19 @@
 sys.path.insert(0, os.path.abspath('../..'))
 
 import sphinx_rtd_theme
+import lightly
 
 
 # -- Project information -----------------------------------------------------
 
 project = 'lightly'
-copyright = '2020, Lightly AG'
+copyright_year = '2020'
+copyright = "Lightly AG"
+website_url = 'https://www.lightly.ai/'
 author = 'Philipp Wirth, Igor Susmelj'
 
 # The full version, including alpha/beta/rc tags
-release = '1.0.0'
+release = lightly.__version__
 master_doc = 'index'
 
 
@@ -79,6 +82,7 @@
 
 html_theme_options = {
     'collapse_navigation': False, # set to false to prevent menu item collapse
+    'logo_only': True
 }
 
 # Add any paths that contain custom static files (such as style sheets) here,
@@ -88,6 +92,10 @@
 
 html_favicon = 'favicon.png'
 
-#html_logo = "../logos/lightly_logo_crop.png"
-def setup(app):
-    app.add_css_file('css/my-styles.css')
+html_logo = '../logos/lightly_logo_crop_white_text.png'
+
+# Exposes variables so that they can be used by django
+html_context = {
+    'copyright_year': copyright_year,
+    'website_url': website_url,
+}
\ No newline at end of file


Test patch:
Repository: lightly-ai/lightly

Issue:
Update swav collate example
The swav collate example uses moco collate by accident:
https://github.com/lightly-ai/lightly/blob/c01ac6cbff2085df604c4fd7f62e77f62e3b05cc/lightly/data/collate.py#L435-L438


Discussion:

Resolving patch:
diff --git a/lightly/data/collate.py b/lightly/data/collate.py
--- a/lightly/data/collate.py
+++ b/lightly/data/collate.py
@@ -433,7 +433,7 @@ class SwaVCollateFunction(MultiCropCollateFunction):
         >>> collate_fn = SwaVCollateFunction()
         >>> 
         >>> # SwaV w/ 2x160 and 4x96 crops 
-        >>> collate_fn = MoCoCollateFunction(
+        >>> collate_fn = SwaVCollateFunction(
         >>>     crop_sizes=[160, 96],
         >>>     crop_counts=[2, 4],
         >>> )


Test patch:
Repository: lightly-ai/lightly

Issue:
Did cifar10 dataset need gaussian blur
In the https://github.com/lightly-ai/lightly/blob/master/lightly/data/collate.py,  class SimCLRCollateFunction() presents a n example for using it,
collate_fn = SimCLRCollateFunction(
        >>>     input_size=32,
        >>>     gaussian_blur=0.,
        >>> )
but in https://docs.lightly.ai/examples/simclr.html
collate_fn = SimCLRCollateFunction(input_size=32)
so I wonder which one is the one you suggested?


Discussion:
Hi @littleolex, in the paper the authors leave out Gaussian blur so I would recommend to use the first option:
```python
collate_fn = SimCLRCollateFunction(
    input_size=32,
    gaussian_blur=0.,
)
```

We will fix the example to do the same 👍 
Thank you for replying so quickly

Resolving patch:
diff --git a/examples/pytorch/simclr.py b/examples/pytorch/simclr.py
--- a/examples/pytorch/simclr.py
+++ b/examples/pytorch/simclr.py
@@ -32,7 +32,10 @@ def forward(self, x):
 # or create a dataset from a folder containing images or videos:
 # dataset = LightlyDataset("path/to/folder")
 
-collate_fn = SimCLRCollateFunction(input_size=32)
+collate_fn = SimCLRCollateFunction(
+    input_size=32,
+    gaussian_blur=0.,
+)
 
 dataloader = torch.utils.data.DataLoader(
     dataset,
diff --git a/examples/pytorch_lightning/simclr.py b/examples/pytorch_lightning/simclr.py
--- a/examples/pytorch_lightning/simclr.py
+++ b/examples/pytorch_lightning/simclr.py
@@ -41,7 +41,10 @@ def configure_optimizers(self):
 # or create a dataset from a folder containing images or videos:
 # dataset = LightlyDataset("path/to/folder")
 
-collate_fn = SimCLRCollateFunction(input_size=32)
+collate_fn = SimCLRCollateFunction(
+    input_size=32,
+    gaussian_blur=0.,
+)
 
 dataloader = torch.utils.data.DataLoader(
     dataset,


Test patch:
Repository: lightly-ai/lightly

Issue:
Incorrect inputsize for BarlowTwins Lightning Example Code
Should the input_size in [1] be `32` instead of `224`?
In [2], we use `input_size=32`.

[1] https://github.com/lightly-ai/lightly/blob/master/examples/pytorch_lightning/barlowtwins.py#L44
[2] https://github.com/lightly-ai/lightly/blob/master/examples/pytorch/barlowtwins.py#L35




Discussion:
Hi @heytitle thanks for the question! I'll pass it on to the author, @guarin 🙂 
I think this is indeed a mistake. 
BarlowTwins works also for small images. 

Note that papers like SwaV use additional smaller crops. Since we already have 32x32 images in cifar10 it makes sense to not use smaller crops there. 
Most of the papers compare benchmarks on the ImageNet dataset with 224x224 input resolution. So if you have larger input images I'd suggest using a larger resolution as well if possible :)

Resolving patch:
diff --git a/examples/pytorch_lightning/barlowtwins.py b/examples/pytorch_lightning/barlowtwins.py
--- a/examples/pytorch_lightning/barlowtwins.py
+++ b/examples/pytorch_lightning/barlowtwins.py
@@ -41,7 +41,7 @@ def configure_optimizers(self):
 # or create a dataset from a folder containing images or videos:
 # dataset = LightlyDataset("path/to/folder")
 
-collate_fn = ImageCollateFunction(input_size=224)
+collate_fn = ImageCollateFunction(input_size=32)
 
 dataloader = torch.utils.data.DataLoader(
     dataset,


Test patch:
Repository: lightly-ai/lightly

Issue:
Unusual error on lightly-1.2.12
Getting this error:

```
Epoch 19: 100% 430/430 [17:42<00:00,  2.47s/it, loss=2.05, v_num=0]
Best model is stored at: /content/lightly_outputs/2022-04-04/12-01-48/lightly_epoch_18.ckpt
########## Starting to embed your dataset.

Error executing job with overrides: ['token=min', 'dataset_id=mine', 'input_dir=/content/drive/MyDrive/data/mine/', 'trainer.max_epochs=20']
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/lightly/cli/lightly_cli.py", line 114, in lightly_cli
    return _lightly_cli(cfg)
  File "/usr/local/lib/python3.7/dist-packages/lightly/cli/lightly_cli.py", line 60, in _lightly_cli
    embeddings = _embed_cli(cfg, is_cli_call)
  File "/usr/local/lib/python3.7/dist-packages/lightly/cli/embed_cli.py", line 83, in _embed_cli
    embeddings, labels, filenames = encoder.embed(dataloader, device=device)
  File "/usr/local/lib/python3.7/dist-packages/lightly/embedding/embedding.py", line 113, in embed
    total=len(dataloader.dataset),
AttributeError: 'BackgroundGenerator' object has no attribute 'dataset'

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
```

There are jpgs in `/content/drive/MyDrive/data/mine/`
Token/dataset_ide correct


Discussion:
Hey @robmarkcole thanks for reporting this bug!

This should be easy to temporarily resolve by uninstalling the `prefetch_generator` package or setting up a new environment without it.

The problem occurs because of this line replacing the PyTorch dataloader by a `BackgroundGenerator`:
https://github.com/lightly-ai/lightly/blob/947d3fdaf5b09988fa491a154470db1e2940c096/lightly/embedding/embedding.py#L110

We'll work on a fix 👍 

Resolving patch:
diff --git a/lightly/embedding/embedding.py b/lightly/embedding/embedding.py
--- a/lightly/embedding/embedding.py
+++ b/lightly/embedding/embedding.py
@@ -106,11 +106,12 @@ def embed(self,
         self.model.eval()
         embeddings, labels, filenames = None, None, []
 
+        dataset = dataloader.dataset
         if lightly._is_prefetch_generator_available():
             dataloader = BackgroundGenerator(dataloader, max_prefetch=3)
         
         pbar = tqdm(
-            total=len(dataloader.dataset),
+            total=len(dataset),
             unit='imgs'
         )
 
@@ -157,7 +158,7 @@ def embed(self,
             embeddings = embeddings.cpu().numpy()
             labels = labels.cpu().numpy()
 
-        sorted_filenames = dataloader.dataset.get_filenames()
+        sorted_filenames = dataset.get_filenames()
         sorted_embeddings = sort_items_by_keys(
             filenames, embeddings, sorted_filenames
         )


Test patch:
Repository: lightly-ai/lightly

Issue:
Loss stuck
Hi, I am trying to run the tutorial posted here 
https://docs.lightly.ai/self-supervised-learning/tutorials/package/tutorial_moco_memory_bank.html
But my loss is stuck at 8.32 after 100 epochs
python 3.9
pytorch-lightning 1.8.1 
lightly 1.2.38

Any suggestions on how I should troubleshoot this?
Thanks in advance!


Discussion:
Hi there,
I'm experiencing the same problem with the latest version of Lightly (1.2.39). However, when I ran the same script with my older version of Lightly, the problem didn't occur and there was no collapse. It looks like the issue was introduced in version 1.2.37 of Lightly (no collapse with 1.2.36), which suggests it may be related to the modification of the MemoryBankModule wich is also used by the Moco loss. I'm also having the same issue when running the [Lightly benchmark on Imagenette](https://docs.lightly.ai/self-supervised-learning/getting_started/benchmarks.html#imagenette) with Moco, with a loss also getting stuck at 8.32 after the first epoch.
Hi, I can confirm that it works on 1.2.36 and not on 1.2.38. 
Looks like we introduced a bug, probably in #993. Will have a look as soon as possible.

Resolving patch:
diff --git a/lightly/loss/memory_bank.py b/lightly/loss/memory_bank.py
--- a/lightly/loss/memory_bank.py
+++ b/lightly/loss/memory_bank.py
@@ -63,7 +63,7 @@ def _init_memory_bank(self, dim: int):
         # https://github.com/facebookresearch/moco but we don't
         # want to pollute our checkpoints
         self.bank = torch.randn(dim, self.size).type_as(self.bank)
-        torch.nn.functional.normalize(self.bank, dim=0)
+        self.bank = torch.nn.functional.normalize(self.bank, dim=0)
         self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)
 
     @torch.no_grad()


Test patch:
Repository: lightly-ai/lightly

Issue:
CUDA errors in NTXentLoss with gloo backend in multi-gpu training 
I was wondering if the `gloo` distributed communication package for multi-gpu training is officially supported by lightly. It seems like e.g. NTXentLoss doesn't work with `gloo` (I'm using pytorch lightning): I get CUDA errors, even when setting `gather_distributed = False`.
I can fix the issue when using `gather_distributed = False` by replacing the line
https://github.com/lightly-ai/lightly/blob/master/lightly/loss/ntx_ent_loss.py#L164
by
```python
labels = labels + batch_size * (dist.rank() if gather_distributed else 0)
```
but then of course I can't use `gather_distributed = True` anymore. 

Using the `nccl` backend, everything works fine, but `nccl` is not working well on some of our machines, so unfortunately I'm stuck with `gloo`. I think using `gloo` might be too much of an exception to fix the problem for `gather_distributed = True`, but maybe it'd help to just replace the line above and mention somewhere in the documentation that `gather_distributed` is not supported for `gloo`?


Discussion:
Hi, gather is not supported for gpus with the `gloo` backend, see: https://pytorch.org/docs/stable/distributed.html#backends
Thanks for pointing that out, I missed that. However, the NTXentLoss still fails even if setting `gather_distributed = False` due to the line mentioned above :)

Resolving patch:
diff --git a/lightly/loss/ntx_ent_loss.py b/lightly/loss/ntx_ent_loss.py
--- a/lightly/loss/ntx_ent_loss.py
+++ b/lightly/loss/ntx_ent_loss.py
@@ -161,7 +161,8 @@ def forward(self, out0: torch.Tensor, out1: torch.Tensor):
 
             # create labels
             labels = torch.arange(batch_size, device=device, dtype=torch.long)
-            labels = labels + dist.rank() * batch_size
+            if self.gather_distributed:
+                labels = labels + dist.rank() * batch_size
             labels = labels.repeat(2)
 
         loss = self.cross_entropy(logits, labels)


Test patch:
Repository: lightly-ai/lightly

Issue:
VICRegL drop default vertical flip
We accidentally pass the horizontal flip probability also as vertical flip probability in `RandomResizedCropAndFlip` transform from `VICRegL`: https://github.com/lightly-ai/lightly/blob/61127338eacd50ad977b9eebaa7a297c404fecb9/lightly/transforms/random_crop_and_flip_with_grid.py#L151-L152

We should fix this by adding vertical flip probability as argument to the class and passing it accordingly to `RandomVerticalFlipWithLocation`. The default vertical flip probability for `VICRegL` should be `0` as only horizontal flip is used.


Discussion:
Hi @guarin can you assign this issue to me. Would like to try out. 
Sure!

Resolving patch:
diff --git a/lightly/transforms/random_crop_and_flip_with_grid.py b/lightly/transforms/random_crop_and_flip_with_grid.py
--- a/lightly/transforms/random_crop_and_flip_with_grid.py
+++ b/lightly/transforms/random_crop_and_flip_with_grid.py
@@ -138,6 +138,7 @@ def __init__(
         crop_min_scale: float = 0.05,
         crop_max_scale: float = 0.2,
         hf_prob: float = 0.5,
+        vf_prob: float = 0.5,
     ):
         super().__init__()
         self.grid_size = grid_size
@@ -145,11 +146,12 @@ def __init__(
         self.crop_min_scale = crop_min_scale
         self.crop_max_scale = crop_max_scale
         self.hf_prob = hf_prob
+        self.vf_prob = vf_prob
         self.resized_crop = RandomResizedCropWithLocation(
             size=self.crop_size, scale=(self.crop_min_scale, self.crop_max_scale)
         )
         self.horizontal_flip = RandomHorizontalFlipWithLocation(self.hf_prob)
-        self.vertical_flip = RandomVerticalFlipWithLocation(self.hf_prob)
+        self.vertical_flip = RandomVerticalFlipWithLocation(self.vf_prob)
 
     def forward(self, img: Image.Image) -> Tuple[Image.Image, torch.Tensor]:
         """Applies random cropping and horizontal flipping to an image, and returns the
diff --git a/lightly/transforms/vicregl_transform.py b/lightly/transforms/vicregl_transform.py
--- a/lightly/transforms/vicregl_transform.py
+++ b/lightly/transforms/vicregl_transform.py
@@ -100,6 +100,7 @@ def __init__(
         global_solarize_prob: float = 0.0,
         local_solarize_prob: float = 0.2,
         hf_prob: float = 0.5,
+        vf_prob: float = 0.0,
         cj_prob: float = 1.0,
         cj_strength: float = 0.5,
         cj_bright: float = 0.8,
@@ -115,6 +116,7 @@ def __init__(
                 crop_min_scale=global_crop_scale[0],
                 crop_max_scale=global_crop_scale[1],
                 hf_prob=hf_prob,
+                vf_prob=vf_prob,
                 grid_size=global_grid_size,
             ),
             VICRegLViewTransform(


Test patch:
Repository: lightly-ai/lightly

Issue:
Remove prefetch generator
## Description

Torch handles prefetching internally and the prefetch generator in PIP is no longer used. We can safely remove it.

## Tasks

- [ ] Remove all code for the prefetch generator
- [ ] Remove dependency for prefetch generator
- [ ] Check if prefetch generator is mentioned in docs


Discussion:

Resolving patch:
diff --git a/lightly/__init__.py b/lightly/__init__.py
--- a/lightly/__init__.py
+++ b/lightly/__init__.py
@@ -94,17 +94,6 @@
     msg = f"Partial import of {__name__}=={__version__} during build process."
     print(msg)
 else:
-    # see if prefetch_generator is available
-    try:
-        import prefetch_generator
-    except ImportError:
-        _prefetch_generator_available = False
-    else:
-        _prefetch_generator_available = True
-
-    def _is_prefetch_generator_available():
-        return _prefetch_generator_available
-
     # see if torchvision vision transformer is available
     try:
         import torchvision.models.vision_transformer
diff --git a/lightly/embedding/embedding.py b/lightly/embedding/embedding.py
--- a/lightly/embedding/embedding.py
+++ b/lightly/embedding/embedding.py
@@ -14,9 +14,6 @@
 from lightly.embedding._base import BaseEmbedding
 from lightly.utils.reordering import sort_items_by_keys
 
-if lightly._is_prefetch_generator_available():
-    from prefetch_generator import BackgroundGenerator
-
 
 class SelfSupervisedEmbedding(BaseEmbedding):
     """Implementation of self-supervised embedding models.
@@ -105,8 +102,6 @@ def embed(
         embeddings, labels, filenames = None, None, []
 
         dataset = dataloader.dataset
-        if lightly._is_prefetch_generator_available():
-            dataloader = BackgroundGenerator(dataloader, max_prefetch=3)
 
         pbar = tqdm(total=len(dataset), unit="imgs")
 


Test patch:
Repository: lightly-ai/lightly

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.



Discussion:

Resolving patch:
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 @@ def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:
             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 @@ def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:
         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
 
 


Test patch:
diff --git a/tests/loss/test_PMSNLoss.py b/tests/loss/test_PMSNLoss.py
--- a/tests/loss/test_PMSNLoss.py
+++ b/tests/loss/test_PMSNLoss.py
@@ -12,7 +12,7 @@
 class TestPMSNLoss:
     def test_regularization_loss(self) -> None:
         criterion = PMSNLoss()
-        mean_anchor_probs = torch.Tensor([0.1, 0.3, 0.6]).log()
+        mean_anchor_probs = torch.Tensor([0.1, 0.3, 0.6])
         loss = criterion.regularization_loss(mean_anchor_probs=mean_anchor_probs)
         norm = 1 / (1**0.25) + 1 / (2**0.25) + 1 / (3**0.25)
         t0 = 1 / (1**0.25) / norm
@@ -45,7 +45,7 @@ def test_forward_cuda(self) -> None:
 class TestPMSNCustomLoss:
     def test_regularization_loss(self) -> None:
         criterion = PMSNCustomLoss(target_distribution=_uniform_distribution)
-        mean_anchor_probs = torch.Tensor([0.1, 0.3, 0.6]).log()
+        mean_anchor_probs = torch.Tensor([0.1, 0.3, 0.6])
         loss = criterion.regularization_loss(mean_anchor_probs=mean_anchor_probs)
         expected_loss = (
             1