Dear @ptrblck,
Thank you so much for your reply.
Assume that in the forward pass I use the following codes to compute output and modify the output:

output = run_function(input)
output = modify(output)
ctx.save_for_backward(output)

where run_function is ReLU, conv., etc. Then, in the backward pass I retrieve the modified output by:

(output, ) = ctx.saved_tensors

Together with the grad_output, will the following code is sufficient to compute the gradient for the input?

torch.autograd.backward(output, grad_output)

If not, what should I do?
Thank you so much for your time.