Refactor MAE to use timm ViT

[guarin]

We want to refactor the MAE code to use the timm ViT backbone instead of torchvision ViT. The reason for it is that torchvision ViT is too different from the ViT used in MAE and doesn't support all the features we need. See #1263 for some issues. Refactoring the code allows use to use the same backbone as was used in the original MAE code.

Ideally the refactoring allows us to be relatively flexible with the backbone we use so that users are free to use different ViT backbones.

Todo

• Refactor MAE code to use timm ViT backbone
• Make sure any timm imports are optional and do not break the package. We want to keep timm dependency optional, see Add timm dependency? #1301

[radiradev]

Hi, I am working on implementing this, should be able to submit a pull request soon.

[guarin]

Hi @radiradev! This issue seems to have gotten lost, the MAE examples are in the meantime already updated to use TIMM ViT. See the docs and the benchmark. Is there something particular that you have found that is still missing? Otherwise I would close this issue.

One thing we didn't yet have time to do is #1367 (I-JEPA with TIMM ViT). Let us know if you would be interested in working on that.

[radiradev]

Hi @guarin, I commented on the wrong issue. (I have started working on #1367).

Feel free to close this one.

[guarin]

Completed in #1461