def apply_masks(x: Tensor, masks: Tensor | list[Tensor]) -> Tensor:
    """Apply masks to the input tensor.

    From https://github.com/facebookresearch/ijepa/blob/main/src/masks/utils.py

    Args:
        x:
            Tensor of shape (B, N, D) where N is the number of patches.
        masks:
            Tensor or list of tensors containing indices of patches in
            [0, N-1] to keep. Each tensor musth have shape (B, K) where K is the number
            of patches to keep. All masks must have the same K.

    Returns:
        Tensor of shape (B * num_masks, K, D) where K is the number of patches to keep.
    """

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

    all_x = []
    for m in masks:
        mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1))
        all_x += [torch.gather(x, dim=1, index=mask_keep)]
    return torch.cat(all_x, dim=0)
Apply masks to the input tensor.

    From https://github.com/facebookresearch/ijepa/blob/main/src/masks/utils.py

    Args:
        x:
            Tensor of shape (B, N, D) where N is the number of patches.
        masks:
            Tensor or list of tensors containing indices of patches in
            [0, N-1] to keep. Each tensor musth have shape (B, K) where K is the number
            of patches to keep. All masks must have the same K.

    Returns:
        Tensor of shape (B * num_masks, K, D) where K is the number of patches to keep.
    
lightly/models/utils.py