def forward(
        self, x: Tensor, step: Optional[int] = None
    ) -> Union[Tensor, List[Tensor]]:
        """Forward pass of the SwaVPrototypes module.

        Args:
            x:
                Input tensor.
            step:
                Current training step.

        Returns:
            The logits after passing through the prototype heads. Returns a single tensor
            if there's one prototype head, otherwise returns a list of tensors.
        """
        self._freeze_prototypes_if_required(step)
        out = []
        for layer in self.heads:
            out.append(layer(x))
        return out[0] if self._is_single_prototype else out
Forward pass of the SwaVPrototypes module.

        Args:
            x:
                Input tensor.
            step:
                Current training step.

        Returns:
            The logits after passing through the prototype heads. Returns a single tensor
            if there's one prototype head, otherwise returns a list of tensors.
        
lightly/models/modules/heads.py