I’m afraid this is not possible with the API that we have for torch::autograd::Function.
There is always the solution of creating the Node by hand but this is going to be more complex (and outside of the official public API).