Convergence of inverse problem of heat equation with PINN

I’m having trouble getting my inverse problem of heat equation with multiple parameters to converge using the neuromancer package (which encodes the PINN model). I am coding using python but was told that I still could ask for feedback here anyway. My code was based on this: Google Colab except I changed the data generation to my own method of data generation and my goal is to recover the parameter a, b and c from the equation dy_t - ady_xx -b dy_x -c y =0.

Through the module deepxde which also uses the PINN model using the same data generation code, I was able to recover the forward solution of this heat equation (still having trouble with the inverse problem) but with neuromancer, I haven’t been able to get convergence with both forward and inverse problem. I used Pytorch backends for both packages.

I am new to machine learning so my issues might be trivial/ overlooked but I would appreciate any pointers! Thanks!