Repository: lightly-ai/lightly (branch: master)
Commit date: 2023-12-08T11:41:32Z

Workflow file: .github/workflows/test_code_format.yml (Code Format Check)
name: Code Format Check

on:
  push:
  pull_request:
  workflow_dispatch:

jobs:
  test:
    name: Check
    runs-on: ubuntu-latest
    services:
      pypi_wayback:
        image: ghcr.io/waleko/pypi-wayback
        ports:
          - 8629:8080
    steps:
    - name: Checkout Code
      uses: actions/checkout@v3
    - name: Hack to get setup-python to work on nektos/act
      run: |
        if [ ! -f "/etc/lsb-release" ] ; then
          echo "DISTRIB_RELEASE=18.04" > /etc/lsb-release
        fi
    - name: Set up Python
      uses: actions/setup-python@v4
      with:
        python-version: "3.7"
    - uses: actions/cache@v2
      with:
        path: ${{ env.pythonLocation }}
        key: cache_v2_${{ env.pythonLocation }}-${{ hashFiles('requirements/**') }}
    - name: Install Dependencies and lightly
      run: pip install -e '.[all]'
    - name: Run Format Check
      run: |
        make format-check
    - name: Run Type Check
      run: |
        make type-check

env:
  PIP_INDEX_URL: http://localhost:8629/2023-12-08
  UV_INDEX_URL: http://localhost:8629/2023-12-08


Failed build logs:
step_name: Check/7_Run Format Check.txt
log: ##[group]Run make format-check
[36;1mmake format-check[0m
shell: /usr/bin/bash -e {0}
env:
  pythonLocation: /opt/hostedtoolcache/Python/3.7.17/x64
  PKG_CONFIG_PATH: /opt/hostedtoolcache/Python/3.7.17/x64/lib/pkgconfig
  Python_ROOT_DIR: /opt/hostedtoolcache/Python/3.7.17/x64
  Python2_ROOT_DIR: /opt/hostedtoolcache/Python/3.7.17/x64
  Python3_ROOT_DIR: /opt/hostedtoolcache/Python/3.7.17/x64
  LD_LIBRARY_PATH: /opt/hostedtoolcache/Python/3.7.17/x64/lib
##[endgroup]
⚫ Checking code format...
isort --check-only --diff .
Skipped 2 files
black --check .
would reformat /home/runner/work/lightly/lightly/lightly/loss/ssley_loss.py

Oh no! 💥 💔 💥
1 file would be reformatted, 343 files would be left unchanged.
make: *** [Makefile:46: format-check] Error 1
##[error]Process completed with exit code 2.


Changed files:
README.md
benchmarks/imagenet/resnet50/main.py
benchmarks/imagenet/resnet50/ssley.py
docs/source/examples/ssley.rst
docs/source/getting_started/benchmarks/imagenette_benchmark.py
docs/source/lightly.loss.rst
examples/pytorch/ssley.py
examples/pytorch_lightning/ssley.py
examples/pytorch_lightning_distributed/ssley.py
lightly/loss/__init__.py
lightly/loss/mmcr_loss.py
lightly/loss/ssley_loss.py
lightly/transforms/__init__.py
lightly/transforms/mmcr_transform.py
lightly/utils/benchmarking/knn_classifier.py
lightly/utils/benchmarking/linear_classifier.py
tests/loss/test_MMCR_loss.py
tests/loss/test_SSLEYLoss.py
tests/transforms/test_mmcr_transform.py
tests/utils/benchmarking/test_knn_classifier.py
tests/utils/benchmarking/test_linear_classifier.py

Fix diff:
diff --git a/README.md b/README.md
index 77e10af3..70afc941 100644
--- a/README.md
+++ b/README.md
@@ -52,7 +52,6 @@ and PyTorch Lightning distributed examples for all models to kickstart your proj
 - SimMIM, 2021 [paper](https://arxiv.org/abs/2111.09886) [docs](https://docs.lightly.ai/self-supervised-learning/examples/simmim.html)
 - SimSiam, 2021 [paper](https://arxiv.org/abs/2011.10566) [docs](https://docs.lightly.ai/self-supervised-learning/examples/simsiam.html)
 - SMoG, 2022 [paper](https://arxiv.org/abs/2207.06167) [docs](https://docs.lightly.ai/self-supervised-learning/examples/smog.html)
-- SSL-EY, 2023 [paper](https://arxiv.org/abs/2310.01012) [docs](https://docs.lightly.ai/self-supervised-learning/examples/ssley.html)
 - SwaV, 2020 [paper](https://arxiv.org/abs/2006.09882) [docs](https://docs.lightly.ai/self-supervised-learning/examples/swav.html)
 - TiCo, 2022 [paper](https://arxiv.org/abs/2206.10698) [docs](https://docs.lightly.ai/self-supervised-learning/examples/tico.html)
 - VICReg, 2022 [paper](https://arxiv.org/abs/2105.04906) [docs](https://docs.lightly.ai/self-supervised-learning/examples/vicreg.html)
@@ -279,17 +278,16 @@ tuned for maximum accuracy. For detailed results and more information about the
 See the [benchmarking scripts](./benchmarks/imagenet/resnet50/) for details.
 
 
-| Model          | Backbone | Batch Size | Epochs | Linear Top1 | Finetune Top1 | kNN Top1 | Tensorboard                                                    | Checkpoint                                                                                                                                                         |
-|----------------|----------|------------|--------|-------------|---------------|----------|----------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| BarlowTwins    | Res50    |        256 |    100 | 62.9        | 72.6          | 45.6     | [link](https://tensorboard.dev/experiment/NxyNRiQsQjWZ82I9b0PvKg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_barlowtwins_2023-08-18_00-11-03/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
-| BYOL           | Res50    |        256 |    100 | 62.4        | 74.0          | 45.6     | [link](https://tensorboard.dev/experiment/Z0iG2JLaTJe5nuBD7DK1bg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_byol_2023-07-10_10-37-32/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt)   |
-| DINO           | Res50    |        128 |    100 | 68.2        | 72.5          | 49.9     | [link](https://tensorboard.dev/experiment/DvKHX9sNSWWqDrRksllPLA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dino_2023-06-06_13-59-48/pretrain/version_0/checkpoints/epoch%3D99-step%3D1000900.ckpt)  |
-| SimCLR*        | Res50    |        256 |    100 | 63.2        | 73.9          | 44.8     | [link](https://tensorboard.dev/experiment/Ugol97adQdezgcVibDYMMA) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_simclr_2023-06-22_09-11-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
-| SimCLR* + DCL  | Res50    |        256 |    100 | 65.1        | 73.5          | 49.6     | [link](https://tensorboard.dev/experiment/k4ZonZ77QzmBkc0lXswQlg/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dcl_2023-07-04_16-51-40/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt)    |
-| SimCLR* + DCLW | Res50    |        256 |    100 | 64.5        | 73.2          | 48.5     | [link](https://tensorboard.dev/experiment/TrALnpwFQ4OkZV3uvaX7wQ/) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dclw_2023-07-07_14-57-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt)   |
-| SSL-EY         | Res50    |        256 |    100 | TODO        | TODO          | TODO     | [link](TODO)| [link](TODO)  |
-| SwAV           | Res50    |        256 |    100 | 67.2        | 75.4          | 49.5     | [link](https://tensorboard.dev/experiment/Ipx4Oxl5Qkqm5Sl5kWyKKg) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_swav_2023-05-25_08-29-14/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt)   |
-| VICReg         | Res50    |        256 |    100 | 63.0        | 73.7          | 46.3     | [link](https://tensorboard.dev/experiment/qH5uywJbTJSzgCEfxc7yUw) | [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_vicreg_2023-09-11_10-53-08/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| Model          | Backbone | Batch Size | Epochs | Linear Top1 | Finetune Top1 | kNN Top1 | Tensorboard | Checkpoint |
+|----------------|----------|------------|--------|-------------|---------------|----------|-------------|------------|
+| BarlowTwins    | Res50    |        256 |    100 |        62.9 |          72.6 |     45.6 |      [link](https://tensorboard.dev/experiment/NxyNRiQsQjWZ82I9b0PvKg/) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_barlowtwins_2023-08-18_00-11-03/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| BYOL           | Res50    |        256 |    100 |        62.4 |          74.0 |     45.6 |      [link](https://tensorboard.dev/experiment/Z0iG2JLaTJe5nuBD7DK1bg) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_byol_2023-07-10_10-37-32/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| DINO           | Res50    |        128 |    100 |        68.2 |          72.5 |     49.9 |      [link](https://tensorboard.dev/experiment/DvKHX9sNSWWqDrRksllPLA) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dino_2023-06-06_13-59-48/pretrain/version_0/checkpoints/epoch%3D99-step%3D1000900.ckpt) |
+| SimCLR*        | Res50    |        256 |    100 |        63.2 |          73.9 |     44.8 |      [link](https://tensorboard.dev/experiment/Ugol97adQdezgcVibDYMMA) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_simclr_2023-06-22_09-11-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| SimCLR* + DCL  | Res50    |        256 |    100 |        65.1 |          73.5 |     49.6 |      [link](https://tensorboard.dev/experiment/k4ZonZ77QzmBkc0lXswQlg/) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dcl_2023-07-04_16-51-40/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| SimCLR* + DCLW | Res50    |        256 |    100 |        64.5 |          73.2 |     48.5 |      [link](https://tensorboard.dev/experiment/TrALnpwFQ4OkZV3uvaX7wQ/) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_dclw_2023-07-07_14-57-13/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| SwAV           | Res50    |        256 |    100 |        67.2 |          75.4 |     49.5 |      [link](https://tensorboard.dev/experiment/Ipx4Oxl5Qkqm5Sl5kWyKKg) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_swav_2023-05-25_08-29-14/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
+| VICReg         | Res50    |        256 |    100 |        63.0 |          73.7 |     46.3 |      [link](https://tensorboard.dev/experiment/qH5uywJbTJSzgCEfxc7yUw) |       [link](https://lightly-ssl-checkpoints.s3.amazonaws.com/imagenet_resnet50_vicreg_2023-09-11_10-53-08/pretrain/version_0/checkpoints/epoch%3D99-step%3D500400.ckpt) |
 
 *\*We use square root learning rate scaling instead of linear scaling as it yields
 better results for smaller batch sizes. See Appendix B.1 in the [SimCLR paper](https://arxiv.org/abs/2002.05709).*
diff --git a/benchmarks/imagenet/resnet50/main.py b/benchmarks/imagenet/resnet50/main.py
index 21906071..8eb3e7f6 100644
--- a/benchmarks/imagenet/resnet50/main.py
+++ b/benchmarks/imagenet/resnet50/main.py
@@ -12,7 +12,6 @@ import finetune_eval
 import knn_eval
 import linear_eval
 import simclr
-import ssley
 import swav
 import torch
 import vicreg
@@ -59,7 +58,6 @@ METHODS = {
     "dclw": {"model": dclw.DCLW, "transform": dclw.transform},
     "dino": {"model": dino.DINO, "transform": dino.transform},
     "simclr": {"model": simclr.SimCLR, "transform": simclr.transform},
-    "ssley": {"model": ssley.SSLEY, "transform": ssley.transform},
     "swav": {"model": swav.SwAV, "transform": swav.transform},
     "vicreg": {"model": vicreg.VICReg, "transform": vicreg.transform},
 }
diff --git a/benchmarks/imagenet/resnet50/ssley.py b/benchmarks/imagenet/resnet50/ssley.py
deleted file mode 100644
index 1c45a75e..00000000
--- a/benchmarks/imagenet/resnet50/ssley.py
+++ /dev/null
@@ -1,127 +0,0 @@
-from typing import List, Tuple
-
-import torch
-from pytorch_lightning import LightningModule
-from torch import Tensor
-from torch.nn import Identity
-from torchvision.models import resnet50
-
-from lightly.loss.ssley_loss import SSLEYLoss
-from lightly.models.modules.heads import VICRegProjectionHead
-from lightly.models.utils import get_weight_decay_parameters
-from lightly.transforms.ssley_transform import VICRegTransform
-from lightly.utils.benchmarking import OnlineLinearClassifier
-from lightly.utils.lars import LARS
-from lightly.utils.scheduler import CosineWarmupScheduler
-
-
-class SSLEY(LightningModule):
-    def __init__(self, batch_size_per_device: int, num_classes: int) -> None:
-        super().__init__()
-        self.save_hyperparameters()
-        self.batch_size_per_device = batch_size_per_device
-
-        resnet = resnet50()
-        resnet.fc = Identity()  # Ignore classification head
-        self.backbone = resnet
-        self.projection_head = VICRegProjectionHead(num_layers=2)
-        self.criterion = SSLEYLoss()
-
-        self.online_classifier = OnlineLinearClassifier(num_classes=num_classes)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.backbone(x)
-
-    def training_step(
-        self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int
-    ) -> Tensor:
-        views, targets = batch[0], batch[1]
-        features = self.forward(torch.cat(views)).flatten(start_dim=1)
-        z = self.projection_head(features)
-        z_a, z_b = z.chunk(len(views))
-        loss = self.criterion(z_a=z_a, z_b=z_b)
-        self.log(
-            "train_loss", loss, prog_bar=True, sync_dist=True, batch_size=len(targets)
-        )
-
-        # Online linear evaluation.
-        cls_loss, cls_log = self.online_classifier.training_step(
-            (features.detach(), targets.repeat(len(views))), batch_idx
-        )
-
-        self.log_dict(cls_log, sync_dist=True, batch_size=len(targets))
-        return loss + cls_loss
-
-    def validation_step(
-        self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int
-    ) -> Tensor:
-        images, targets = batch[0], batch[1]
-        features = self.forward(images).flatten(start_dim=1)
-        cls_loss, cls_log = self.online_classifier.validation_step(
-            (features.detach(), targets), batch_idx
-        )
-        self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets))
-        return cls_loss
-
-    def configure_optimizers(self):
-        # Don't use weight decay for batch norm, bias parameters, and classification
-        # head to improve performance.
-        params, params_no_weight_decay = get_weight_decay_parameters(
-            [self.backbone, self.projection_head]
-        )
-        global_batch_size = self.batch_size_per_device * self.trainer.world_size
-        base_lr = _get_base_learning_rate(global_batch_size=global_batch_size)
-        optimizer = LARS(
-            [
-                {"name": "ssley", "params": params},
-                {
-                    "name": "ssley_no_weight_decay",
-                    "params": params_no_weight_decay,
-                    "weight_decay": 0.0,
-                },
-                {
-                    "name": "online_classifier",
-                    "params": self.online_classifier.parameters(),
-                    "weight_decay": 0.0,
-                },
-            ],
-            # Linear learning rate scaling with a base learning rate of 0.2.
-            # See https://arxiv.org/pdf/2105.04906.pdf for details.
-            lr=base_lr * global_batch_size / 256,
-            momentum=0.9,
-            weight_decay=1e-6,
-        )
-        scheduler = {
-            "scheduler": CosineWarmupScheduler(
-                optimizer=optimizer,
-                warmup_epochs=(
-                    self.trainer.estimated_stepping_batches
-                    / self.trainer.max_epochs
-                    * 10
-                ),
-                max_epochs=self.trainer.estimated_stepping_batches,
-                end_value=0.01,  # Scale base learning rate from 0.2 to 0.002.
-            ),
-            "interval": "step",
-        }
-        return [optimizer], [scheduler]
-
-
-# SSLEY transform
-transform = VICRegTransform()
-
-
-def _get_base_learning_rate(global_batch_size: int) -> float:
-    """Returns the base learning rate for training 100 epochs with a given batch size.
-
-    This follows section C.4 in https://arxiv.org/pdf/2105.04906.pdf.
-
-    """
-    if global_batch_size == 128:
-        return 0.8
-    elif global_batch_size == 256:
-        return 0.5
-    elif global_batch_size == 512:
-        return 0.4
-    else:
-        return 0.3
diff --git a/docs/source/examples/ssley.rst b/docs/source/examples/ssley.rst
deleted file mode 100644
index 8a449f69..00000000
--- a/docs/source/examples/ssley.rst
+++ /dev/null
@@ -1,48 +0,0 @@
-.. _ssley:
-
-SSL-EY
-=======
-
-SSL-EY is a method that explicitly
-avoids the collapse problem with a simple regularization term on the variance of the embeddings along each dimension individually. It inherits the model structure from
-`Barlow Twins, 2022 <https://arxiv.org/abs/2103.03230>`_ changing the loss. Doing so allows the stabilization of the training and leads to performance improvements. 
-
-Reference:
-    `Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2023 <https://arxiv.org/abs/2310.01012>`_
-
-
-.. tabs::
-    .. tab:: PyTorch
-
-        This example can be run from the command line with::
-
-            python lightly/examples/pytorch/ssley.py
-
-        .. literalinclude:: ../../../examples/pytorch/ssley.py
-
-    .. tab:: Lightning
-
-        This example can be run from the command line with::
-
-            python lightly/examples/pytorch_lightning/ssley.py
-
-        .. literalinclude:: ../../../examples/pytorch_lightning/ssley.py
-
-    .. tab:: Lightning Distributed
-
-        This example runs on multiple gpus using Distributed Data Parallel (DDP)
-        training with Pytorch Lightning. At least one GPU must be available on 
-        the system. The example can be run from the command line with::
-
-            python lightly/examples/pytorch_lightning_distributed/ssley.py
-
-        The model differs in the following ways from the non-distributed
-        implementation:
-
-        - Distributed Data Parallel is enabled
-        - Distributed Sampling is used in the dataloader
-
-        Distributed Sampling makes sure that each distributed process sees only
-        a subset of the data.
-
-        .. literalinclude:: ../../../examples/pytorch_lightning_distributed/ssley.py
\ No newline at end of file
diff --git a/docs/source/getting_started/benchmarks/imagenette_benchmark.py b/docs/source/getting_started/benchmarks/imagenette_benchmark.py
index af2759ed..b5c5e8e7 100644
--- a/docs/source/getting_started/benchmarks/imagenette_benchmark.py
+++ b/docs/source/getting_started/benchmarks/imagenette_benchmark.py
@@ -29,7 +29,6 @@ Results (4.5.2023):
 | SimCLR           |        256 |    200 |                  0.835 |   49.7 Min |      3.7 GByte |
 | SimMIM (ViT-B32) |        256 |    200 |                  0.315 |  115.5 Min |      9.7 GByte |
 | SimSiam          |        256 |    200 |                  0.752 |   58.2 Min |      3.9 GByte |
-| SSL-EY           |        256 |    200 |                  TO-DO |   TO-DO    |    TO-DO GByte |
 | SwaV             |        256 |    200 |                  0.861 |   73.3 Min |      6.4 GByte |
 | SwaVQueue        |        256 |    200 |                  0.827 |   72.6 Min |      6.4 GByte |
 | SMoG             |        256 |    200 |                  0.663 |   58.7 Min |      2.6 GByte |
@@ -51,7 +50,6 @@ Results (4.5.2023):
 | SimCLR           |        256 |    800 |                  0.889 |  193.5 Min |      3.7 GByte |
 | SimMIM (ViT-B32) |        256 |    800 |                  0.343 |  446.5 Min |      9.7 GByte |
 | SimSiam          |        256 |    800 |                  0.872 |  206.4 Min |      3.9 GByte |
-| SSL-EY           |        256 |    800 |                  TO-DO |   TO-DO    |    TO-DO GByte |
 | SwaV             |        256 |    800 |                  0.902 |  283.2 Min |      6.4 GByte |
 | SwaVQueue        |        256 |    800 |                  0.890 |  282.7 Min |      6.4 GByte |
 | SMoG             |        256 |    800 |                  0.788 |  232.1 Min |      2.6 GByte |
@@ -83,7 +81,6 @@ from lightly.loss import (
     NegativeCosineSimilarity,
     NTXentLoss,
     PMSNLoss,
-    SSLEYLoss,
     SwaVLoss,
     TiCoLoss,
     VICRegLLoss,
@@ -270,7 +267,6 @@ def create_dataset_train_ssl(model):
         SimCLRModel: simclr_transform,
         SimMIMModel: simmim_transform,
         SimSiamModel: simsiam_transform,
-        SSL_EYModel: vicreg_transform,
         SwaVModel: swav_transform,
         SwaVQueueModel: swav_transform,
         SMoGModel: smog_transform,
@@ -1169,42 +1165,6 @@ class SimMIMModel(BenchmarkModule):
         return [optim], [cosine_scheduler]
 
 
-class SSLEYModel(BenchmarkModule):
-    def __init__(self, dataloader_kNN, num_classes):
-        super().__init__(dataloader_kNN, num_classes)
-        # create a ResNet backbone and remove the classification head
-        resnet = torchvision.models.resnet18()
-        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
-        self.projection_head = heads.BarlowTwinsProjectionHead(512, 2048, 2048)
-        self.criterion = SSLEYLoss()
-        self.warmup_epochs = 40 if max_epochs >= 800 else 20
-
-    def forward(self, x):
-        x = self.backbone(x).flatten(start_dim=1)
-        z = self.projection_head(x)
-        return z
-
-    def training_step(self, batch, batch_index):
-        (x0, x1), _, _ = batch
-        z0 = self.forward(x0)
-        z1 = self.forward(x1)
-        loss = self.criterion(z0, z1)
-        return loss
-
-    def configure_optimizers(self):
-        # Training diverges without LARS
-        optim = LARS(
-            self.parameters(),
-            lr=0.3 * lr_factor,
-            weight_decay=1e-4,
-            momentum=0.9,
-        )
-        cosine_scheduler = scheduler.CosineWarmupScheduler(
-            optim, self.warmup_epochs, max_epochs
-        )
-        return [optim], [cosine_scheduler]
-
-
 class VICRegModel(BenchmarkModule):
     def __init__(self, dataloader_kNN, num_classes):
         super().__init__(dataloader_kNN, num_classes)
@@ -1450,7 +1410,6 @@ models = [
     SimCLRModel,
     # SimMIMModel, # disabled by default because SimMIM uses larger images with size 224
     SimSiamModel,
-    SSLEYModel,
     SwaVModel,
     SwaVQueueModel,
     SMoGModel,
diff --git a/docs/source/lightly.loss.rst b/docs/source/lightly.loss.rst
index 2976cca5..38e304b6 100644
--- a/docs/source/lightly.loss.rst
+++ b/docs/source/lightly.loss.rst
@@ -41,9 +41,6 @@ lightly.loss
 .. autoclass:: lightly.loss.regularizer.co2.CO2Regularizer
    :members:
 
-.. autoclass:: lightly.loss.ssley_loss.SSLEYLoss
-   :members:
-
 .. autoclass:: lightly.loss.swav_loss.SwaVLoss
    :members:
 
diff --git a/examples/pytorch/ssley.py b/examples/pytorch/ssley.py
deleted file mode 100644
index 695fe92d..00000000
--- a/examples/pytorch/ssley.py
+++ /dev/null
@@ -1,70 +0,0 @@
-import torch
-import torchvision
-from torch import nn
-
-## The projection head is the same as the Barlow Twins one
-from lightly.loss import SSLEYLoss
-
-## The projection head is the same as the Barlow Twins one
-from lightly.loss.ssley_loss import SSLEYLoss
-from lightly.models.modules.heads import VICRegProjectionHead
-from lightly.transforms.vicreg_transform import VICRegTransform
-
-
-class SSLEY(nn.Module):
-    def __init__(self, backbone):
-        super().__init__()
-        self.backbone = backbone
-        self.projection_head = VICRegProjectionHead(
-            input_dim=512,
-            hidden_dim=2048,
-            output_dim=2048,
-            num_layers=2,
-        )
-
-    def forward(self, x):
-        x = self.backbone(x).flatten(start_dim=1)
-        z = self.projection_head(x)
-        return z
-
-
-resnet = torchvision.models.resnet18()
-backbone = nn.Sequential(*list(resnet.children())[:-1])
-model = SSLEY(backbone)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-transform = VICRegTransform(input_size=32)
-dataset = torchvision.datasets.CIFAR10(
-    "datasets/cifar10", download=True, transform=transform
-)
-# or create a dataset from a folder containing images or videos:
-# dataset = LightlyDataset("path/to/folder", transform=transform)
-
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=256,
-    shuffle=True,
-    drop_last=True,
-    num_workers=8,
-)
-criterion = SSLEYLoss()
-optimizer = torch.optim.SGD(model.parameters(), lr=0.06)
-
-print("Starting Training")
-for epoch in range(10):
-    total_loss = 0
-    for batch in dataloader:
-        x0, x1 = batch[0]
-        x0 = x0.to(device)
-        x1 = x1.to(device)
-        z0 = model(x0)
-        z1 = model(x1)
-        loss = criterion(z0, z1)
-        total_loss += loss.detach()
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
-    avg_loss = total_loss / len(dataloader)
-    print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}")
diff --git a/examples/pytorch_lightning/ssley.py b/examples/pytorch_lightning/ssley.py
deleted file mode 100644
index 6cc17de9..00000000
--- a/examples/pytorch_lightning/ssley.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# Note: The model and training settings do not follow the reference settings
-# from the paper. The settings are chosen such that the example can easily be
-# run on a small dataset with a single GPU.
-
-import pytorch_lightning as pl
-import torch
-import torchvision
-from torch import nn
-
-from lightly.loss.ssley_loss import SSLEYLoss
-
-## The projection head is the same as the Barlow Twins one
-from lightly.models.modules.heads import VICRegProjectionHead
-from lightly.transforms.vicreg_transform import VICRegTransform
-
-
-class SSLEY(pl.LightningModule):
-    def __init__(self):
-        super().__init__()
-        resnet = torchvision.models.resnet18()
-        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
-        self.projection_head = VICRegProjectionHead(
-            input_dim=512,
-            hidden_dim=2048,
-            output_dim=2048,
-            num_layers=2,
-        )
-        self.criterion = SSLEYLoss()
-
-    def forward(self, x):
-        x = self.backbone(x).flatten(start_dim=1)
-        z = self.projection_head(x)
-        return z
-
-    def training_step(self, batch, batch_index):
-        (x0, x1) = batch[0]
-        z0 = self.forward(x0)
-        z1 = self.forward(x1)
-        loss = self.criterion(z0, z1)
-        return loss
-
-    def configure_optimizers(self):
-        optim = torch.optim.SGD(self.parameters(), lr=0.06)
-        return optim
-
-
-model = SSLEY()
-
-transform = VICRegTransform(input_size=32)
-dataset = torchvision.datasets.CIFAR10(
-    "datasets/cifar10", download=True, transform=transform
-)
-# or create a dataset from a folder containing images or videos:
-# dataset = LightlyDataset("path/to/folder", transform=transform)
-
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=256,
-    shuffle=True,
-    drop_last=True,
-    num_workers=8,
-)
-
-accelerator = "gpu" if torch.cuda.is_available() else "cpu"
-
-trainer = pl.Trainer(max_epochs=10, devices=1, accelerator=accelerator)
-trainer.fit(model=model, train_dataloaders=dataloader)
diff --git a/examples/pytorch_lightning_distributed/ssley.py b/examples/pytorch_lightning_distributed/ssley.py
deleted file mode 100644
index 3f034adc..00000000
--- a/examples/pytorch_lightning_distributed/ssley.py
+++ /dev/null
@@ -1,77 +0,0 @@
-# Note: The model and training settings do not follow the reference settings
-# from the paper. The settings are chosen such that the example can easily be
-# run on a small dataset with a single GPU.
-
-import pytorch_lightning as pl
-import torch
-import torchvision
-from torch import nn
-
-from lightly.loss import SSLEYLoss
-
-## The projection head is the same as the Barlow Twins one
-from lightly.models.modules.heads import VICRegProjectionHead
-from lightly.transforms.vicreg_transform import VICRegTransform
-
-
-class SSLEY(pl.LightningModule):
-    def __init__(self):
-        super().__init__()
-        resnet = torchvision.models.resnet18()
-        self.backbone = nn.Sequential(*list(resnet.children())[:-1])
-        self.projection_head = VICRegProjectionHead(
-            input_dim=512,
-            hidden_dim=2048,
-            output_dim=2048,
-            num_layers=2,
-        )
-
-        # enable gather_distributed to gather features from all gpus
-        # before calculating the loss
-        self.criterion = SSLEYLoss(gather_distributed=True)
-
-    def forward(self, x):
-        x = self.backbone(x).flatten(start_dim=1)
-        z = self.projection_head(x)
-        return z
-
-    def training_step(self, batch, batch_index):
-        (x0, x1) = batch[0]
-        z0 = self.forward(x0)
-        z1 = self.forward(x1)
-        loss = self.criterion(z0, z1)
-        return loss
-
-    def configure_optimizers(self):
-        optim = torch.optim.SGD(self.parameters(), lr=0.06)
-        return optim
-
-
-model = SSLEY()
-
-transform = VICRegTransform(input_size=32)
-dataset = torchvision.datasets.CIFAR10(
-    "datasets/cifar10", download=True, transform=transform
-)
-# or create a dataset from a folder containing images or videos:
-# dataset = LightlyDataset("path/to/folder", transform=transform)
-
-dataloader = torch.utils.data.DataLoader(
-    dataset,
-    batch_size=256,
-    shuffle=True,
-    drop_last=True,
-    num_workers=8,
-)
-
-# Train with DDP and use Synchronized Batch Norm for a more accurate batch norm
-# calculation. Distributed sampling is also enabled with replace_sampler_ddp=True.
-trainer = pl.Trainer(
-    max_epochs=10,
-    devices="auto",
-    accelerator="gpu",
-    strategy="ddp",
-    sync_batchnorm=True,
-    use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
-)
-trainer.fit(model=model, train_dataloaders=dataloader)
diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py
index 51ff573c..385cf3bb 100644
--- a/lightly/loss/__init__.py
+++ b/lightly/loss/__init__.py
@@ -5,11 +5,11 @@
 from lightly.loss.barlow_twins_loss import BarlowTwinsLoss
 from lightly.loss.dcl_loss import DCLLoss, DCLWLoss
 from lightly.loss.dino_loss import DINOLoss
+from lightly.loss.mmcr_loss import MMCRLoss
 from lightly.loss.msn_loss import MSNLoss
 from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity
 from lightly.loss.ntx_ent_loss import NTXentLoss
 from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss
-from lightly.loss.ssley_loss import SSLEYLoss
 from lightly.loss.swav_loss import SwaVLoss
 from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss
 from lightly.loss.tico_loss import TiCoLoss
diff --git a/lightly/loss/mmcr_loss.py b/lightly/loss/mmcr_loss.py
new file mode 100644
index 00000000..883b1147
--- /dev/null
+++ b/lightly/loss/mmcr_loss.py
@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+from torch.linalg import svd
+
+
+class MMCRLoss(nn.Module):
+    """Implementation of the loss function from MMCR [0] using Manifold Capacity.
+    All hyperparameters are set to the default values from the paper for ImageNet.
+
+    - [0]: Efficient Coding of Natural Images using Maximum Manifold Capacity
+            Representations, 2023, https://arxiv.org/pdf/2303.03307.pdf
+
+     Examples:
+
+
+        >>> # initialize loss function
+        >>> loss_fn = MMCRLoss()
+        >>> transform = MMCRTransform(k=2)
+        >>>
+        >>> # transform images, then feed through encoder and projector
+        >>> x = transform(x)
+        >>> online = online_network(x)
+        >>> momentum = momentum_network(x)
+        >>>
+        >>> # calculate loss
+        >>> loss = loss_fn(online, momentum)
+    """
+
+    def __init__(self, lmda: float = 5e-3):
+        super().__init__()
+        if lmda < 0:
+            raise ValueError("lmda must be greater than or equal to 0")
+
+        self.lmda = lmda
+
+    def forward(self, online: torch.Tensor, momentum: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            online:
+                Output of the online network for the current batch. Expected to be
+                of shape (batch_size, k, embedding_size), where k represents the
+                number of randomly augmented views for each sample.
+            momentum:
+                Output of the momentum network for the current batch. Expected to be
+                of shape (batch_size, k, embedding_size), where k represents the
+                number of randomly augmented views for each sample.
+
+        """
+        assert (
+            online.shape == momentum.shape
+        ), "online and momentum need to have the same shape"
+
+        B = online.shape[0]
+
+        # Concatenate and calculate centroid
+        z = torch.cat([online, momentum], dim=1)
+        c = torch.mean(z, dim=1)  # B x D
+
+        # Calculate singular values
+        _, S_z, _ = svd(z)
+        _, S_c, _ = svd(c)
+
+        # Calculate loss
+        loss = -1.0 * torch.sum(S_c) + self.lmda * torch.sum(S_z) / B
+
+        return loss
diff --git a/lightly/loss/ssley_loss.py b/lightly/loss/ssley_loss.py
deleted file mode 100644
index d7f015d6..00000000
--- a/lightly/loss/ssley_loss.py
+++ /dev/null
@@ -1,83 +0,0 @@
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-from torch import Tensor
-
-from lightly.utils.dist import gather
-
-
-class SSLEYLoss(torch.nn.Module):
-    """Implementation of the SSL-EY loss [0].
-
-    - [0] Efficient Algorithms for the CCA Family: Unconstrained Objectives with Unbiased Gradients, 2022, https://arxiv.org/abs/2310.01012
-
-    Attributes:
-        gather_distributed:
-            If True then the cross-correlation matrices from all gpus are gathered and
-            summed before the loss calculation.
-        eps:
-            Epsilon for numerical stability.
-
-    Examples:
-
-        >>> # initialize loss function
-        >>> loss_fn = SSLEYLoss()
-        >>>
-        >>> # generate two random transforms of images
-        >>> t0 = transforms(images)
-        >>> t1 = transforms(images)
-        >>>
-        >>> # feed through model
-        >>> out0, out1 = model(t0, t1)
-        >>>
-        >>> # calculate loss
-        >>> loss = loss_fn(out0, out1)
-    """
-
-    def __init__(
-        self,
-        gather_distributed: bool = False,
-        eps=0.0001,
-    ):
-        super(SSLEYLoss, self).__init__()
-        if gather_distributed and not dist.is_available():
-            raise ValueError(
-                "gather_distributed is True but torch.distributed is not available. "
-                "Please set gather_distributed=False or install a torch version with "
-                "distributed support."
-            )
-
-        self.gather_distributed = gather_distributed
-        self.eps = eps
-
-    def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
-        """Returns SSL-EY loss.
-
-        Args:
-            z_a:
-                Tensor with shape (batch_size, ..., dim).
-            z_b:
-                Tensor with shape (batch_size, ..., dim).
-        """
-        assert (
-                z_a.shape[0] > 1 and z_b.shape[0] > 1
-        ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}"
-        assert (
-                z_a.shape == z_b.shape
-        ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}."
-        # gather all batches
-        if self.gather_distributed and dist.is_initialized():
-            world_size = dist.get_world_size()
-            if world_size > 1:
-                z_a = torch.cat(gather(z_a), dim=0)
-                z_b = torch.cat(gather(z_b), dim=0)
-
-        z_a = z_a - z_a.mean(dim=0)
-        z_b = z_b - z_b.mean(dim=0)
-        batch_size = z_a.size(0)
-        C = 2*(z_a.T @ z_b) / (batch_size - 1)
-        V = (z_a.T @ z_a) / (batch_size - 1) + (z_b.T @ z_b) / (batch_size - 1)
-
-        loss = -2 * torch.trace(C) + torch.trace(V @ V)
-
-        return loss
diff --git a/lightly/transforms/__init__.py b/lightly/transforms/__init__.py
index 33c70bda..1f841fcc 100644
--- a/lightly/transforms/__init__.py
+++ b/lightly/transforms/__init__.py
@@ -18,6 +18,7 @@ from lightly.transforms.fast_siam_transform import FastSiamTransform
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.jigsaw import Jigsaw
 from lightly.transforms.mae_transform import MAETransform
+from lightly.transforms.mmcr_transform import MMCRTransform
 from lightly.transforms.moco_transform import MoCoV1Transform, MoCoV2Transform
 from lightly.transforms.msn_transform import MSNTransform, MSNViewTransform
 from lightly.transforms.pirl_transform import PIRLTransform
diff --git a/lightly/transforms/mmcr_transform.py b/lightly/transforms/mmcr_transform.py
new file mode 100644
index 00000000..4d47a3e2
--- /dev/null
+++ b/lightly/transforms/mmcr_transform.py
@@ -0,0 +1,89 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+from lightly.transforms.byol_transform import BYOLView1Transform
+from lightly.transforms.multi_view_transform import MultiViewTransform
+from lightly.transforms.utils import IMAGENET_NORMALIZE
+
+
+class MMCRTransform(MultiViewTransform):
+    """Implements the transformations for MMCR[0], which
+    are based on BYOL[1].
+
+    Input to this transform:
+        PIL Image or Tensor.
+
+    Output of this transform:
+        List of Tensor of length k.
+
+    Applies the following augmentations by default:
+        - Random resized crop
+        - Random horizontal flip
+        - Color jitter
+        - Random gray scale
+        - Gaussian blur
+        - Solarization
+        - ImageNet normalization
+
+    Please refer to the BYOL implementation for additional details.
+
+    - [0]: Efficient Coding of Natural Images using Maximum Manifold Capacity
+            Representations, 2023, https://arxiv.org/pdf/2303.03307.pdf
+    - [1]: Bootstrap Your Own Latent, 2020, https://arxiv.org/pdf/2006.07733.pdf
+
+
+    Input to this transform:
+        PIL Image or Tensor.
+
+    Output of this transform:
+        List of tensors of length k.
+
+    Attributes:
+        k: Number of views.
+        transform: The transform to apply to each view.
+    """
+
+    def __init__(
+        self,
+        k: int = 8,
+        input_size: int = 224,
+        cj_prob: float = 0.8,
+        cj_strength: float = 1.0,
+        cj_bright: float = 0.4,
+        cj_contrast: float = 0.4,
+        cj_sat: float = 0.2,
+        cj_hue: float = 0.1,
+        min_scale: float = 0.08,
+        random_gray_scale: float = 0.2,
+        gaussian_blur: float = 1.0,
+        solarization_prob: float = 0.0,
+        kernel_size: Optional[float] = None,
+        sigmas: Tuple[float, float] = (0.1, 2),
+        vf_prob: float = 0.0,
+        hf_prob: float = 0.5,
+        rr_prob: float = 0.0,
+        rr_degrees: Union[None, float, Tuple[float, float]] = None,
+        normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE,
+    ):
+        if k < 1:
+            raise ValueError("k must be greater than or equal to 1")
+        transform = BYOLView1Transform(
+            input_size=input_size,
+            cj_prob=cj_prob,
+            cj_strength=cj_strength,
+            cj_bright=cj_bright,
+            cj_contrast=cj_contrast,
+            cj_sat=cj_sat,
+            cj_hue=cj_hue,
+            min_scale=min_scale,
+            random_gray_scale=random_gray_scale,
+            gaussian_blur=gaussian_blur,
+            solarization_prob=solarization_prob,
+            kernel_size=kernel_size,
+            sigmas=sigmas,
+            vf_prob=vf_prob,
+            hf_prob=hf_prob,
+            rr_prob=rr_prob,
+            rr_degrees=rr_degrees,
+            normalize=normalize,
+        )
+        super().__init__(transforms=[transform] * k)
diff --git a/lightly/utils/benchmarking/knn_classifier.py b/lightly/utils/benchmarking/knn_classifier.py
index 6b5b7a95..ab78bb8c 100644
--- a/lightly/utils/benchmarking/knn_classifier.py
+++ b/lightly/utils/benchmarking/knn_classifier.py
@@ -6,7 +6,6 @@ from pytorch_lightning import LightningModule
 from torch import Tensor
 from torch.nn import Module
 
-from lightly.models.utils import activate_requires_grad, deactivate_requires_grad
 from lightly.utils.benchmarking import knn_predict
 from lightly.utils.benchmarking.topk import mean_topk_accuracy
 
@@ -97,6 +96,7 @@ class KNNClassifier(LightningModule):
         self._train_features_tensor: Union[Tensor, None] = None
         self._train_targets_tensor: Union[Tensor, None] = None
 
+    @torch.no_grad()
     def training_step(self, batch, batch_idx) -> None:
         images, targets = batch[0], batch[1]
         features = self.model.forward(images).flatten(start_dim=1)
@@ -141,13 +141,9 @@ class KNNClassifier(LightningModule):
             targets = targets.flatten().t().contiguous()
             self._train_targets_tensor = targets.to(self.device)
 
-    def on_fit_start(self) -> None:
-        # Freeze model weights.
-        deactivate_requires_grad(model=self.model)
-
-    def on_fit_end(self) -> None:
-        # Unfreeze model weights.
-        activate_requires_grad(model=self.model)
+    def on_train_epoch_start(self) -> None:
+        # Set model to eval mode to disable norm layer updates.
+        self.model.eval()
 
     def configure_optimizers(self) -> None:
         # configure_optimizers must be implemented for PyTorch Lightning. Returning None
diff --git a/lightly/utils/benchmarking/linear_classifier.py b/lightly/utils/benchmarking/linear_classifier.py
index 3fb76c6e..647b8dd4 100644
--- a/lightly/utils/benchmarking/linear_classifier.py
+++ b/lightly/utils/benchmarking/linear_classifier.py
@@ -1,11 +1,11 @@
 from typing import Any, Dict, List, Tuple, Union
 
+import torch
 from pytorch_lightning import LightningModule
 from torch import Tensor
 from torch.nn import CrossEntropyLoss, Linear, Module
 from torch.optim import SGD, Optimizer
 
-from lightly.models.utils import activate_requires_grad, deactivate_requires_grad
 from lightly.utils.benchmarking.topk import mean_topk_accuracy
 from lightly.utils.scheduler import CosineWarmupScheduler
 
@@ -93,7 +93,11 @@ class LinearClassifier(LightningModule):
         self.criterion = CrossEntropyLoss()
 
     def forward(self, images: Tensor) -> Tensor:
-        features = self.model.forward(images).flatten(start_dim=1)
+        if self.freeze_model:
+            with torch.no_grad():
+                features = self.model.forward(images).flatten(start_dim=1)
+        else:
+            features = self.model.forward(images).flatten(start_dim=1)
         output: Tensor = self.classification_head(features)
         return output
 
@@ -147,12 +151,7 @@ class LinearClassifier(LightningModule):
         }
         return [optimizer], [scheduler]
 
-    def on_fit_start(self) -> None:
-        # Freeze model weights.
-        if self.freeze_model:
-            deactivate_requires_grad(model=self.model)
-
-    def on_fit_end(self) -> None:
-        # Unfreeze model weights.
+    def on_train_epoch_start(self) -> None:
         if self.freeze_model:
-            activate_requires_grad(model=self.model)
+            # Set model to eval mode to disable norm layer updates.
+            self.model.eval()
diff --git a/tests/loss/test_MMCR_loss.py b/tests/loss/test_MMCR_loss.py
new file mode 100644
index 00000000..b0437ba1
--- /dev/null
+++ b/tests/loss/test_MMCR_loss.py
@@ -0,0 +1,65 @@
+import unittest
+
+import torch
+
+from lightly.loss.mmcr_loss import MMCRLoss
+
+
+class testMMCRLoss(unittest.TestCase):
+    def test_forward(self) -> None:
+        bs = 3
+        dim = 128
+        k = 32
+
+        loss_fn = MMCRLoss()
+        online = torch.randn(bs, k, dim)
+        momentum = torch.randn(bs, k, dim)
+
+        loss = loss_fn(online, momentum)
+
+        print(loss)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "cuda not available")
+    def test_forward_cuda(self) -> None:
+        bs = 3
+        dim = 128
+        k = 32
+
+        loss_fn = MMCRLoss()
+        online = torch.randn(bs, k, dim).cuda()
+        momentum = torch.randn(bs, k, dim).cuda()
+
+        loss = loss_fn(online, momentum)
+
+        print(loss)
+
+    def test_loss_value(self) -> None:
+        """If all values are zero, the loss should be zero."""
+        bs = 3
+        dim = 128
+        k = 32
+
+        loss_fn = MMCRLoss()
+        online = torch.zeros(bs, k, dim)
+        momentum = torch.zeros(bs, k, dim)
+
+        loss = loss_fn(online, momentum)
+
+        self.assertTrue(loss == 0)
+
+    def test_lambda_value_error(self) -> None:
+        """If lambda is negative, a ValueError should be raised."""
+        with self.assertRaises(ValueError):
+            MMCRLoss(lmda=-1)
+
+    def test_shape_assertion_forward(self) -> None:
+        bs = 3
+        dim = 128
+        k = 32
+
+        loss_fn = MMCRLoss()
+        online = torch.randn(bs, k, dim)
+        momentum = torch.randn(bs, k, dim + 1)
+
+        with self.assertRaises(AssertionError):
+            loss_fn(online, momentum)
diff --git a/tests/loss/test_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py
deleted file mode 100644
index 0ee7b1e2..00000000
--- a/tests/loss/test_SSLEYLoss.py
+++ /dev/null
@@ -1,65 +0,0 @@
-import unittest
-
-import pytest
-import torch
-from pytest_mock import MockerFixture
-from torch import distributed as dist
-
-from lightly.loss import SSLEYLoss
-
-
-class TestSSLEYLoss:
-    def test__gather_distributed(self, mocker: MockerFixture) -> None:
-        mock_is_available = mocker.patch.object(dist, "is_available", return_value=True)
-        SSLEYLoss(gather_distributed=True)
-        mock_is_available.assert_called_once()
-
-    def test__gather_distributed_dist_not_available(
-        self, mocker: MockerFixture
-    ) -> None:
-        mock_is_available = mocker.patch.object(
-            dist, "is_available", return_value=False
-        )
-        with pytest.raises(ValueError):
-            SSLEYLoss(gather_distributed=True)
-        mock_is_available.assert_called_once()
-
-
-class TestSSLEYLossUnitTest(unittest.TestCase):
-    # Old tests in unittest style, please add new tests to TestSSLEYLoss using pytest.
-    def test_forward_pass(self):
-        loss = SSLEYLoss()
-        for bsz in range(2, 4):
-            x0 = torch.randn((bsz, 32))
-            x1 = torch.randn((bsz, 32))
-
-            # symmetry
-            l1 = loss(x0, x1)
-            l2 = loss(x1, x0)
-            self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0)
-
-    @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available")
-    def test_forward_pass_cuda(self):
-        loss = SSLEYLoss()
-        for bsz in range(2, 4):
-            x0 = torch.randn((bsz, 32)).cuda()
-            x1 = torch.randn((bsz, 32)).cuda()
-
-            # symmetry
-            l1 = loss(x0, x1)
-            l2 = loss(x1, x0)
-            self.assertAlmostEqual((l1 - l2).pow(2).item(), 0.0)
-
-    def test_forward_pass__error_batch_size_1(self):
-        loss = SSLEYLoss()
-        x0 = torch.randn((1, 32))
-        x1 = torch.randn((1, 32))
-        with self.assertRaises(AssertionError):
-            loss(x0, x1)
-
-    def test_forward_pass__error_different_shapes(self):
-        loss = SSLEYLoss()
-        x0 = torch.randn((2, 32))
-        x1 = torch.randn((2, 16))
-        with self.assertRaises(AssertionError):
-            loss(x0, x1)
diff --git a/tests/transforms/test_mmcr_transform.py b/tests/transforms/test_mmcr_transform.py
new file mode 100644
index 00000000..c51c9e62
--- /dev/null
+++ b/tests/transforms/test_mmcr_transform.py
@@ -0,0 +1,21 @@
+import pytest
+from PIL import Image
+
+from lightly.transforms.mmcr_transform import MMCRTransform
+
+
+def test_raise_value_error() -> None:
+    with pytest.raises(ValueError):
+        MMCRTransform(k=0)
+
+
+def test_num_views() -> None:
+    multi_view_transform = MMCRTransform(k=3)
+    assert len(multi_view_transform.transforms) == 3
+
+
+def test_multi_view_on_pil_image() -> None:
+    multi_view_transform = MMCRTransform(k=3)
+    sample = Image.new("RGB", (100, 100))
+    output = multi_view_transform(sample)
+    assert len(output) == 3
diff --git a/tests/utils/benchmarking/test_knn_classifier.py b/tests/utils/benchmarking/test_knn_classifier.py
index e8ccfc4f..1b776ce7 100644
--- a/tests/utils/benchmarking/test_knn_classifier.py
+++ b/tests/utils/benchmarking/test_knn_classifier.py
@@ -61,7 +61,11 @@ class TestKNNClassifier:
 
     def _test__accelerator(self, accelerator: str, expected_device: str) -> None:
         torch.manual_seed(0)
-        model = nn.Linear(3, 2)
+        linear = nn.Linear(3, 2)
+        batch_norm = nn.BatchNorm1d(2)
+        model = nn.Sequential(linear, batch_norm)
+        initial_weights = linear.weight.clone()
+        initial_bn_weights = batch_norm.weight.clone()
         classifier = KNNClassifier(model, num_classes=10, knn_k=20)
         trainer = Trainer(max_epochs=1, accelerator=accelerator, devices=1)
         train_features = torch.randn(40, 3)
@@ -93,6 +97,16 @@ class TestKNNClassifier:
         assert classifier._train_features_tensor.device.type == expected_device
         assert classifier._train_targets_tensor.device.type == expected_device
 
+        # Verify that model weights were not updated.
+        assert torch.all(torch.eq(initial_weights, linear.weight))
+        assert torch.all(torch.eq(initial_bn_weights, batch_norm.weight))
+        # Verify that batch norm statistics were not updated. Note that even though the
+        # running mean was not updated, it is still initialized to zero.
+        assert batch_norm.running_mean is not None
+        assert torch.all(
+            torch.eq(batch_norm.running_mean, torch.zeros_like(batch_norm.running_mean))
+        )
+
     def test__features_dtype(self) -> None:
         model = nn.Identity()
         # Set feature_dtype to torch.int to test if classifier correctly changes dtype.
diff --git a/tests/utils/benchmarking/test_linear_classifier.py b/tests/utils/benchmarking/test_linear_classifier.py
index 6118733d..14d2e2ef 100644
--- a/tests/utils/benchmarking/test_linear_classifier.py
+++ b/tests/utils/benchmarking/test_linear_classifier.py
@@ -68,8 +68,10 @@ class TestLinearClassifier:
         train_dataloader = DataLoader(dataset, batch_size=2)
         val_dataloader = DataLoader(dataset, batch_size=2)
         linear = nn.Linear(3 * 8 * 8, 4)
-        model = nn.Sequential(nn.Flatten(), linear)
+        batch_norm = nn.BatchNorm1d(4)
+        model = nn.Sequential(nn.Flatten(), linear, batch_norm)
         initial_weights = linear.weight.clone()
+        initial_bn_weights = batch_norm.weight.clone()
         linear_classifier = LinearClassifier(
             model=model,
             batch_size_per_device=2,
@@ -97,6 +99,13 @@ class TestLinearClassifier:
 
         # Verify that model weights were not updated.
         assert torch.all(torch.eq(initial_weights, linear.weight))
+        assert torch.all(torch.eq(initial_bn_weights, batch_norm.weight))
+        # Verify that batch norm statistics were not updated. Note that even though the
+        # running mean was not updated, it is still initialized to zero.
+        assert batch_norm.running_mean is not None
+        assert torch.all(
+            torch.eq(batch_norm.running_mean, torch.zeros_like(batch_norm.running_mean))
+        )
         # Verify that head weights were updated.
         assert not torch.all(
             torch.eq(initial_head_weights, linear_classifier.classification_head.weight)