This helped me a lot! Thank you, Omer.
I am also translating a physics informed neural network (PINN) from Tensorflow (1.0) to PyTorch and struggled with some syntaxes as well.
At the moment my script (based on this notebook 5 by Perdikaris) is running in PyTorch and gives similar results, however I still have some questions regarding my coding. My script is waaayyyy slower than Tensorflow (Tf trains in about 1 minute, but Torch in 40 minutes(!!) with the same amount of iterations and datapoints). Have you encountered this problem as well? I figured out that adding the create_graphs in torch.autograd.grad(dx.sum(), x, create_graph=True) slows down the program. Of course, without this, the program won’t run properly. Also the behaviour of the losses is a bit different as can be seen in the pictures (left is Tensorflow, right PyTorch).
Loss_comparison1200×500 101 KB
Do you have any idea how to upgrade the PyTorch code such that the training time will be similar to the training time of the Tensorflow model? If required, I could show (some parts of) the code.