Hi,
First, given that this is the last piece of your network, you don’t have to wrap it into a custom autograd.Function, you can also just compute the loss and gradient as a regular python functions and then do:
output = model(input)
loss, grad = my_loss(output, target)
opt.zero_grad()
output.backward(grad)
opt.step()

If you really want to use a custom Function because it is cleaner, I don’t think i can see anything wrong with this one.
Are you sure that the gradients it compute as grad_input are properly non-zero?