def _freeze_prototypes_if_required(self, step: Optional[int] = None) -> None:
        """Freezes the prototypes if the specified number of steps has been reached."""
        if self.n_steps_frozen_prototypes > 0:
            if step is None:
                raise ValueError(
                    "`n_steps_frozen_prototypes` is greater than 0, please"
                    " provide the `step` argument to the `forward()` method."
                )
            self.requires_grad_(step >= self.n_steps_frozen_prototypes)
Freezes the prototypes if the specified number of steps has been reached.
lightly/models/modules/heads.py