def forward(self, x, masks_x, masks):
        """Forward pass of the IJEPAPredictor.

        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"

        if not isinstance(masks_x, list):
            masks_x = [masks_x]

        if not isinstance(masks, list):
            masks = [masks]

        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)

        x = self.ln(self.layers(x))

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

        return x
Forward pass of the IJEPAPredictor.

        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.py