def __init__(
        self,
        input_dim: int = 128,
        n_prototypes: Union[List[int], int] = 3000,
        n_steps_frozen_prototypes: int = 0,
    ):
        """Intializes the SwaVPrototypes module with the specified parameters"""
        super(SwaVPrototypes, self).__init__()

        # Default to a list of 1 if n_prototypes is an int.
        self.n_prototypes = (
            n_prototypes if isinstance(n_prototypes, list) else [n_prototypes]
        )
        self._is_single_prototype = True if isinstance(n_prototypes, int) else False
        self.heads = nn.ModuleList(
            [nn.Linear(input_dim, prototypes) for prototypes in self.n_prototypes]
        )
        self.n_steps_frozen_prototypes = n_steps_frozen_prototypes
Intializes the SwaVPrototypes module with the specified parameters
lightly/models/modules/heads.py