Gradient update is extremely slow

[AshStuff]

I am trying to train a Image-Text Contrastive learning model and I am using a Functional Approach. The number of grad steps are 32 and the batch size per step is 32 which makes the total batch size as 1024.

Any idea how to increase the speed while doing the gradient update ?

[liuweie]

I am trying to train a Image-Text Contrastive learning model and I am using a Functional Approach. The number of grad steps are 32 and the batch size per step is 32 which makes the total batch size as 1024.

Any idea how to increase the speed while doing the gradient update ?

Hello , I meet same question as you , did you figure it out?

[liyongkang123]

Hi all, I am facing the same issue: the speed is very slow. I also observed that my GPU memory was only a few GB in use, even though I had increased the batch size. I want to ask how you handled this problem. Thank you so much for your help!