How to switch between gpu and cpu when backprop?

Hello. I have a problem because of running out of gpu memory. In my nn module, I’ve define a custom functional.However, this functional is too memory-demanding. Then I wonder if we can switch to cpu mode when the program is up to run my custom functional. I think I can invoke the .cpu() function and them pass the feature map to the custom functional, however what will happen when backprop the gradient, can pytorch auto tranfer the gradient from gpu to cpu in the custom functional and tranfer from cpu to gpu after going out of the custom functional? Thanks!

Hi,

Yes, if you call .cpu() on the variable before giving it to your function, it will be converted back to whatever was the type before the call when doing the backward.