I’m afraid there is no simple way to do this today.
We are working on adding this though: SavedVariable default hooks · Issue #58659 · pytorch/pytorch · GitHub 18 (this should be done in 1-2 months).
The workaround with using custom Function would work but you will need to save on the ctx foo.cpu() and not use save_for_backward. And then in the backward formula do ctx.foo.cuda().
But that means that you have to rewrite most of the backward formulas which is very painful.