Cuda out of memory when training ViT-B architecture on ImageNet 224x224 with a batch size of 1024

[jiachenlei]

Hi, Thanks for your great work and open-sourcing your codes!

I am working on one of my project that utilizes contrastive loss. My code computes the contrastive loss for three times each with different input. I follow the instructions in README.md and simply replace code snippets that compute contrastive loss with inf_cl_loss. However, on ImageNet 224x224, when training a ViT-B model with a batch size of 1024 utilizing 8xA800, the out-of-memory exception occurred.

In addition to this repo, I also utilized the xformers repository for efficient ViT attention computation.

Do you have any clues why this is happening? Thanks!