Hi,
If you’re using nightly or in the upcoming 1.10, you will be able to register such hooks directly: Autograd mechanics — PyTorch master documentation 1
And if you just want to move everything from GPU to CPU, we have a builtin here: Automatic differentiation package - torch.autograd — PyTorch master documentation 1