Thanks Tom, I’ll have a try. So does it means I could just make the additional output (like the out2 in your example) non differentiable and add two more position in the python code of backward(), without changing the CUDA kernel? Thanks!