def forward(
        self,
        x: Tensor,
        masks_x: list[Tensor] | Tensor,
        masks: list[Tensor] | Tensor,
    ) -> Tensor:
        """Forward pass of the IJEPAPredictorTIMM.

        Args:
            x:
                Input tensor.
            masks_x:
                Mask indices for the input tensor.
            masks:
                Mask indices for the predicted tokens.

        Returns:
            The predicted output tensor.
        """

        assert (masks is not None) and (
            masks_x is not None
        ), "Cannot run predictor without mask indices"

        len_masks_x = len(masks_x) if isinstance(masks_x, list) else 1
        len_masks = len(masks) if isinstance(masks, list) else 1

        B = len(x) // len_masks_x
        x = self.predictor_embed(x)
        x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1)

        x += utils.apply_masks(x_pos_embed, masks_x)
        _, N_ctxt, _ = x.shape

        pos_embs = self.predictor_pos_embed.repeat(B, 1, 1)
        pos_embs = utils.apply_masks(pos_embs, masks)
        pos_embs = utils.repeat_interleave_batch(pos_embs, B, repeat=len_masks_x)
        pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1)

        pred_tokens += pos_embs
        x = x.repeat(len_masks, 1, 1)
        x = torch.cat([x, pred_tokens], dim=1)

        for blk in self.predictor_blocks:
            x = blk(x)
        x = self.predictor_norm(x)

        x = x[:, N_ctxt:]
        x = self.predictor_proj(x)

        return x
Forward pass of the IJEPAPredictorTIMM.

        Args:
            x:
                Input tensor.
            masks_x:
                Mask indices for the input tensor.
            masks:
                Mask indices for the predicted tokens.

        Returns:
            The predicted output tensor.
        
lightly/models/modules/ijepa_timm.py