Hi, I need to construct a loss function that uses the gradient of the output of the NN w.r.t the input, i.e. something like:

1 output = NN(input, weight)
2 grad_input = output.backward(retain_variables=True)
3 loss = criterion(grad_input.grad.data, target_of_grad_input)
4 loss.backward()

The problem is that after step 2, I can only have the numerical value of the gradient of output w.r.t to input (i.e. grad_input). Then it is no longer a function of the weight, and as a result the loss is not a function of weight. Then step 4 (where I want to compute the loss w.r.t the weight) will fail.
How should I tackle problem like this? Thanks for suggestions.