I’m trying to transfer memory between GPU and CPU with gradient checkpointing, i use this code https://pytorch.org/docs/stable/_modules/torch/utils/checkpoint.html 7
In my forward function i tried to use:
tensor = args[0]
tensor = tensor.cpu()
tensor = tensor.cuda()

When i transfer the memory to the cpu, _grad_fn of tensor is set to None.
Does anyone know how to avoid this?