Hi, I have a question related to autograd. After finishing FP, I want to transfer the input tensor from GPU to CPU, and then after a while transfer it back to GPU for backward and obtain grads. The problem is that if I transfer the tensor back, it won’t go back to the computation graph, and therefore unable to obtain grads from the tensor.
I know there is way that I hook the original GPU tensor to obtain grad. But I wonder if it’s possible to obtain the grad in the transferring back stage.
Here is a simple example:
import torch

device = torch.device("cuda:0")
cpu = torch.device("cpu")

a = torch.randn(5, requires_grad=True, device=device)
b = torch.randn(5, requires_grad=True, device=device)

with torch.autograd.graph.save_on_cpu():
    d = a * b

grads = []

def hook_wrapper(grads):
    def hook(input_gradient):
        grads.append(input_gradient)

    return hook

a = a.to(cpu)
a = a.to(device)
a.register_hook(hook_wrapper(grads))

d.sum().backward()

print(len(grads))
print(grads)

where the len(grads) will be 0 and I cannot obtain the grad.