Fix broken IJEPA Example

[MalteEbner]

When running python examples/pytorch/ijepa.py, the code fails:

  File "/Users/malteebnerlightly/Documents/GitHub/lightly/lightly/models/modules/ijepa.py", line 226, in forward
    input = utils.apply_masks(input, idx_keep)

The utils are imported as from lightly.models import utils.

The function lightly.models/utils.py:apply_masks() was added in a55cf44 and removed later.

[vectorvp]

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

[guarin]

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)

[vectorvp]

@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?

[vectorvp]

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.