@shubham_vats You can either use .retain_grad() on the input value but a more consistent way of getting this value is to use hooks. You can use a full backward hook (not a backward hook as that’s deprecated and gives the wrong result). Also, check there’s no use of in-place operations (for example with ReLU) nor use of .detach() which breaks the gradient of your computational graph.
@Avani_Gupta For an explanation of what grad_input and grad_output mean there’s an explanation on the forums here. The grad_output term is the gradient of the loss with respect to the output of that nn.Module.