So I didn’t want to be one of “those guys” who say nvm I solved the problem but in this case my I think that it was a lot of small bugs in my own implementation and I didn’t think anyone could benefit from it.
That said, the most important change I made was to change
        dxx, = torch.autograd.grad(dx.sum(), x, retain_graph=True) to
        dxx, = torch.autograd.grad(dx.sum(), x, create_graph=True)
This is important because otherwise dxx is computed but not added to the graph i.e. it is treated as a constant when doing the backprob (derivation wrt weights). Without the dxx term Burgers equation reduces to the inviscid Burgers equation which is what happend to me!
Other than that if I learned one thing while working on this is that training PINNs (I’m guessing you are woking on something similar) is extremely tricky.
PD: as the second derivative of ReLU is 0, when your equation has a dxx somewhere it usually fails to learn it! Took me a while to figure that one out.
Hope this helps.
Edit: So I can’t send you a pm because I’m still a new user u.u I’m still working on this problem and I’ll be happy to talk about it if you want. I also have the working code for this (beware, it is not clean!) that I can share.