Custom Loss function, loss is not improving

[ptrblck]

There might be a few issues with your loss calculation:

• you are detaching the outputs from the computation graph by assigning them to the argmax. The softmax function might be an alternative.
  To quote @KFrank:

@KFrank

Consider the softmax function (which should more properly be
called the “soft-argmax” function): It can be understood as a
differentiable approximation to the argmax function (which is
not differentiable because it jumps around discretely).

• If you are using non-PyTorch methods, e.g. numpy methods, you will also break the computation graph and would need to write the backward function manually. However, if you can swap all numpy methods for their PyTorch counterparts, Autograd should create the backward automatically.

• Rewrapping a tensor will also break the graph: R_opt = Variable(R_opt.data, requires_grad=True). Also the usage of the .data attribute is not recommended, as Autograd cannot track these changes, which might result in a wrong gradient calculation. Variables are deprecated since PyTorch 0.4., so you can use tensors in newer versions.