You hardly ever need one hot labels, I don’t think you need them here. Also, I would use torch.nn.functional.cross_entropy_loss instead of re-instantiating the module over and over. You do need predictions, but these are per-class encoded anyways.
You would want to use log_softmax rather than softmax + log.
For the loss involving a gradient w.r.t. the inputs, you’d need to use use requires_grad_() on the inputs and then use gr_input, = torch.autograd.grad(sum_log_hat_y, input)  to get the gradient of your loss computation.
If you believe your input gradients to be reasonably independent of the other inputs in the batch (this can be problematic for BatchNorm but is fine for most other things), you can sum over the batch before taking the gradient, this will make things much faster because you only need one gradient.
I think you are missing a **2 before summing.

Best regards
Thomas