Hi,



 M_S:

Is there a problem with using differential equations as loss terms when training NNs? Or maybe my way of calculating the derivatives is not desirable?


There is no issue with using derivatives in loss term. For instance, you can solve Poisson PDE using its equation as loss function (and applying boundary/initial condition).
I think there might be an issue with your model definition or the way you obtain grads.
Siren paper might be able to help you, please see this notebook 4.
It solves PDEs by including the equation in loss function and/or supervising derivative of ground truth.
As @tom pointed precisely, best way is to use torch.autograd.grad. For instance,
def gradient(y, x, grad_outputs=None):
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]
    return grad

Bests