How many memory do I need for traning a kmax model when I use the resnet50 as the backbone?

[Janet-H]

I use the Resnet50 as the backbone in 2 GPUs(A100 80G per GPU) , but it always say OOM. So I reduce the batch_size and crop , when the batchsize is 32 and cropsize is 321 the kmax model can train. However the PQ is far away from yours , so can you tell me how many memory do I need for training the same result as yours.