Repository: lightly-ai/lightly (branch: master)
Commit date: 2023-12-06T13:39:07Z

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-06
  UV_INDEX_URL: http://localhost:8629/2023-12-06


Failed build logs:
step_name: Check/8_Run Type Check.txt
log: ##[group]Run make type-check
[36;1mmake type-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]
mypy lightly tests
lightly/loss/ssley_loss.py:6: note: In module imported here:
lightly/utils/dist.py: note: In member "forward" of class "GatherLayer":
lightly/utils/dist.py:16: error: Function is missing a type annotation for one or more arguments  [no-untyped-def]
lightly/utils/dist.py: note: In class "GatherLayer":
lightly/utils/dist.py:16: error: Type of decorated function contains type "Any" ("Callable[[Any, Tensor], tuple[Tensor, ...]]")  [misc]
lightly/utils/dist.py:16: error: Signature of "forward" incompatible with supertype "Function"  [override]
lightly/utils/dist.py:16: note:      Superclass:
lightly/utils/dist.py:16: note:          @staticmethod
lightly/utils/dist.py:16: note:          def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any
lightly/utils/dist.py:16: note:      Subclass:
lightly/utils/dist.py:16: note:          @staticmethod
lightly/utils/dist.py:16: note:          def forward(ctx: Any, input: Tensor) -> tuple[Tensor, ...]
lightly/utils/dist.py: note: In member "backward" of class "GatherLayer":
lightly/utils/dist.py:23: error: Function is missing a type annotation for one or more arguments  [no-untyped-def]
lightly/utils/dist.py: note: In class "GatherLayer":
lightly/utils/dist.py:23: error: Type of decorated function contains type "Any" ("Callable[[Any, VarArg(Tensor)], Tensor]")  [misc]
lightly/utils/dist.py: note: In function "gather":
lightly/utils/dist.py:42: error: Returning Any from function declared to return "tuple[Tensor]"  [no-any-return]
lightly/utils/dist.py: note: In function "rank_zero_only":
lightly/utils/dist.py:73: error: Function is missing a type annotation  [no-untyped-def]
lightly/utils/dist.py:85: error: Function is missing a type annotation  [no-untyped-def]
lightly/utils/dist.py: note: At top level:
lightly/utils/dist.py:92: error: Untyped decorator makes function "print_rank_zero" untyped  [misc]
lightly/utils/dist.py: note: In function "print_rank_zero":
lightly/utils/dist.py:93: error: Function is missing a type annotation for one or more arguments  [no-untyped-def]
lightly/utils/dist.py: note: At top level:
lightly/utils/dist.py:93: error: Function is untyped after decorator transformation  [misc]
tests/loss/test_SSLEYLoss.py:10: error: Module "lightly.loss" has no attribute "SSLEYLoss"  [attr-defined]
tests/loss/test_SSLEYLoss.py: note: In member "test_forward_pass" of class "TestSSLEYLossUnitTest":
tests/loss/test_SSLEYLoss.py:32: error: Function is missing a return type annotation  [no-untyped-def]
tests/loss/test_SSLEYLoss.py:32: note: Use "-> None" if function does not return a value
tests/loss/test_SSLEYLoss.py: note: In member "test_forward_pass_cuda" of class "TestSSLEYLossUnitTest":
tests/loss/test_SSLEYLoss.py:44: error: Function is missing a return type annotation  [no-untyped-def]
tests/loss/test_SSLEYLoss.py:44: note: Use "-> None" if function does not return a value
tests/loss/test_SSLEYLoss.py: note: In class "TestSSLEYLossUnitTest":
tests/loss/test_SSLEYLoss.py:44: error: Type of decorated function contains type "Any" ("Callable[[TestSSLEYLossUnitTest], Any]")  [misc]
tests/loss/test_SSLEYLoss.py: note: In member "test_forward_pass__error_batch_size_1" of class "TestSSLEYLossUnitTest":
tests/loss/test_SSLEYLoss.py:55: error: Function is missing a return type annotation  [no-untyped-def]
tests/loss/test_SSLEYLoss.py:55: note: Use "-> None" if function does not return a value
tests/loss/test_SSLEYLoss.py: note: In member "test_forward_pass__error_different_shapes" of class "TestSSLEYLossUnitTest":
tests/loss/test_SSLEYLoss.py:62: error: Function is missing a return type annotation  [no-untyped-def]
tests/loss/test_SSLEYLoss.py:62: note: Use "-> None" if function does not return a value
lightly/loss/ssley_loss.py: note: In member "__init__" of class "SSLEYLoss":
lightly/loss/ssley_loss.py:37: error: Function is missing a type annotation for one or more arguments  [no-untyped-def]
lightly/loss/ssley_loss.py:50: error: Name "lambda_param" is not defined  [name-defined]
lightly/loss/ssley_loss.py:51: error: Name "mu_param" is not defined  [name-defined]
lightly/loss/ssley_loss.py:52: error: Name "nu_param" is not defined  [name-defined]
lightly/loss/ssley_loss.py: note: In member "forward" of class "SSLEYLoss":
lightly/loss/ssley_loss.py:92: error: Returning Any from function declared to return "Tensor"  [no-any-return]
Found 22 errors in 3 files (checked 346 source files)
make: *** [Makefile:68: type-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/ssley_loss.py
lightly/utils/benchmarking/knn_classifier.py
lightly/utils/benchmarking/linear_classifier.py
tests/loss/test_SSLEYLoss.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/ssley_loss.py b/lightly/loss/ssley_loss.py
deleted file mode 100644
index 74dc1e1e..00000000
--- a/lightly/loss/ssley_loss.py
+++ /dev/null
@@ -1,140 +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.lambda_param = lambda_param
-        self.mu_param = mu_param
-        self.nu_param = nu_param
-        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}."
-
-        # invariance term of the loss
-        inv_loss = invariance_loss(x=z_a, y=z_b)
-
-        # 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)
-
-        var_loss = 0.5 * (
-            variance_loss(x=z_a, eps=self.eps) + variance_loss(x=z_b, eps=self.eps)
-        )
-        cov_loss = covariance_loss(x=z_a) + covariance_loss(x=z_b)
-
-        loss = (
-            self.lambda_param * inv_loss
-            + self.mu_param * var_loss
-            + self.nu_param * cov_loss
-        )
-        return loss
-
-
-def invariance_loss(x: Tensor, y: Tensor) -> Tensor:
-    """Returns SSL-EY invariance loss.
-
-    Args:
-        x:
-            Tensor with shape (batch_size, ..., dim).
-        y:
-            Tensor with shape (batch_size, ..., dim).
-    """
-    return F.mse_loss(x, y)
-
-
-def variance_loss(x: Tensor, eps: float = 0.0001) -> Tensor:
-    """Returns SSL-EY variance loss.
-
-    Args:
-        x:
-            Tensor with shape (batch_size, ..., dim).
-        eps:
-            Epsilon for numerical stability.
-    """
-    x = x - x.mean(dim=0)
-    std = torch.sqrt(x.var(dim=0) + eps)
-    loss = torch.mean(F.relu(1.0 - std))
-    return loss
-
-
-def covariance_loss(x: Tensor) -> Tensor:
-    """Returns SSL-EY covariance loss.
-
-    Generalized version of the covariance loss with support for tensors with more than
-    two dimensions.
-
-    Args:
-        x:
-            Tensor with shape (batch_size, ..., dim).
-    """
-    x = x - x.mean(dim=0)
-    batch_size = x.size(0)
-    dim = x.size(-1)
-    # nondiag_mask has shape (dim, dim) with 1s on all non-diagonal entries.
-    nondiag_mask = ~torch.eye(dim, device=x.device, dtype=torch.bool)
-    # cov has shape (..., dim, dim)
-    cov = torch.einsum("b...c,b...d->...cd", x, x) / (batch_size - 1)
-    loss = cov[..., nondiag_mask].pow(2).sum(-1) / dim
-    return loss.mean()
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_SSLEYLoss.py b/tests/loss/test_SSLEYLoss.py
deleted file mode 100644
index a803e588..00000000
--- a/tests/loss/test_SSLEYLoss.py
+++ /dev/null
@@ -1,67 +0,0 @@
-import unittest
-
-import pytest
-import torch
-import torch.nn.functional as F
-from pytest_mock import MockerFixture
-from torch import Tensor
-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/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)