Hi all,
I converted a tensorflow code (link to tf code) to PyTorch and it all works fine (results are comparable in my opinion). However, the code in PyTorch is way slower than in TensorFlow. In TensorFlow, the training only takes 1 minute, where the PyTorch trains in over 40 minutes (I can’t use cuda on my laptop). It uses the same amount of iterations, optimizer, loss function, etc. My question is, do you know anything which might slow down the PyTorch code this much and is there anything I could do about it?
A brief overview of what I did. If you require the code-snippet I could post that too, of course.
I literally translated the tf code to torch, but I do not have much experience in PyTorch.

Everything (loss, optimizer, initialization, training, forward pass,…) is still written in a (nn) class, as is done in the tf code, although I read it is not convenient to do that in PyTorch.
the loss function is written as a specific function of the class.
I wrote down the training ‘session’ as:

self.train_op.zero_grad()
loss = self.loss()
loss.backward(retain_graph = True)
self.train_op.step()


The derivatives of the differential equation are given as (in which the sigma the normalization factor is):

u_x = torch.autograd.grad(outputs = u.sum(), inputs = x, create_graph = True)[0]/self.sigma_x
u_xx = torch.autograd.grad(outputs = u_x.sum(), inputs = x, create_graph = True)[0]/self.sigma_x

I can understand that the ‘create_graph=True’ option slows down the training, but not as much as 40 times.
Do you have any general advice on making the PyTorch code faster?