Hi,
You can easily move Tensors from cpu to gpu as these operations are differentiable in pytorch.
The problem though is that you won’t be able to move the buffers contained in the computational graph: there is no common interface on how to store these for each Function and so it will be special cases for each