Hi I am having some trouble figuring out why my code (in particular the loss function) runs faster on a CPU (~ 0.03 sec) than on a GPU (~ 0.2 sec).
Is there a way I can fix this?
Here is how I define my loss:
def loss(X, u_r):
t1 = perf_counter() ################################# For testing
X.requires_grad = True
u_r_X = torch.zeros(X.shape[0], X.shape[1], X.shape[1], device=device)
for i in range(X.shape[0]):
u_r_X[i,:,:] = torch.autograd.functional.jacobian(u_r, X[i][None,:], create_graph=True)[0,:,0,:]
energy_tensor = torch.zeros(X.shape[0], 1, device=device)
for j in range(X.shape[0]):
F = torch.eye(u_r_X[j].shape[0], device=device) + u_r_X[j]
energy_tensor[j] = 0.5 * (torch.sum((F)**2) - 2.0) - torch.log((torch.det(F))) + 50.0 * torch.log((torch.det(F)))**2
t2 = perf_counter() ################################# For testing
print(t2-t1, ' secs') ################################# For testing
return torch.mean(energy_tensor)