Hi,
I’m a bit confused in the sense that the gradient flowing back depend on the value that was forwarded. So if the last value is X or Z here, the gradient will be different. Is that fine?
If it is, then you can simply compute the gradient s for X first and then backprop them through Z with Z.backward(grad_X).