def __init__(
        self,
        num_patches: int,
        patch_size: int,
        in_chans: int = 3,
        embed_dim: int = 1024,
        decoder_embed_dim: int = 512,
        decoder_depth: int = 8,
        decoder_num_heads: int = 16,
        mlp_ratio: float = 4.0,
        proj_drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        norm_layer: Callable[..., nn.Module] = partial(LayerNorm, eps=1e-6),
        initialize_weights: bool = True,
        mask_token: Optional[Parameter] = None,
    ):
        """Initializes the MAEDecoderTIMM with the specified parameters."""

        super().__init__()

        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
        self.mask_token = (
            nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
            if mask_token is None
            else mask_token
        )

        # Positional encoding of the decoder
        self.decoder_pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False
        )  # fixed sin-cos embedding

        self.decoder_blocks = Sequential(
            *[
                Block(
                    decoder_embed_dim,
                    decoder_num_heads,
                    mlp_ratio,
                    qkv_bias=True,
                    norm_layer=norm_layer,
                    proj_drop=proj_drop_rate,
                    attn_drop=attn_drop_rate,
                )
                for i in range(decoder_depth)
            ]
        )

        self.decoder_norm = norm_layer(decoder_embed_dim)
        self.decoder_pred = nn.Linear(
            decoder_embed_dim, patch_size**2 * in_chans, bias=True
        )  # decoder to patch

        if initialize_weights:
            self._initialize_weights()
Initializes the MAEDecoderTIMM with the specified parameters.
lightly/models/modules/masked_autoencoder_timm.py