Hi,
You want gradients to flow back all the way to the vectors that were given as input to the __init__ function?
If so, your code will work fine.
But if you want it to work with .to()-like functions, it is harder. Because then you want these to be done in a differentiable manner, which they don’t really do. And this will create part of the graph that will be re-used between forward passes and that can cause your backward to fail because some buffers have been freed.