Strange CUDA: OOM (Out of Memory) Issue

[Vishawjeet-rmsl]

Hi, So I'm currently trying to train SAM with different configurations.
So I've modified the depth and patch size in the codes, I'm facing one strange issue.
Below is the vit-h (ignore attention indexes...)

and Below is the vit-b:

This is the patch size (kept same for both)

To my surprise, the model with vit-h is training fine:

But the model with vit-b keeps giving me OOM Issue:

I'm confused, how and why! Vit-B is lighter than Vit-H, if Vit-H is getting trained, then Vit-B should train as well.

BTW, I'm also training the Image encoder in both cases.