@MalteEbner, hello, so the idea is to return apply_masks() to utils.py?

Yes that would be a great start. I'm not sure if something else is broken as well. The relevant code from the original PR is:

def apply_masks(x, masks):
    """
    :param x: tensor of shape [B (batch-size), N (num-patches), D (feature-dim)]
    :param masks: list of tensors containing indices of patches in [N] to keep
    """
    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)

@guarin, apply_masks together with repeat_interleave_batch were moved to lightly/models/modules/ijepa_timm.py. However, timm is used nowhere. Does this module make sense? Do we want to remove it and just add the functions we need to utils.py?

I've delved into the conversations from past issues regarding this timm module and plans to move to it from the torchvision. Based on that I just moved apply_masks together with repeat_interleave_batch to utils.py from lightly/models/modules/ijepa_timm.py.

Now the example works. PR is attached.