Pytorch backprop медленнее по сравнению с Tensorflow?

Я реализовал простую сеть DDQN в pytorch и tensorflow. Сеть довольно мелкая. Хотя прямой проход намного быстрее в PyTorch по сравнению с TF, шаг обратного распространения намного медленнее по сравнению с TF. Оба шага обратного распространения были выполнены на ЦП. Есть идеи, как его улучшить.

Сетевая часть:

def __init__(self, hidden_size_IP=100, hidden_size_rest=100, alpha=0.01, state_size=27, action_size=8,
             learning_rate=1e-6):
    super().__init__()

    # build hidden layers
    self.l1 = nn.Sequential(nn.Linear(in_features=500, out_features=400),
                            nn.LeakyReLU(negative_slope=alpha))
    self.l2 = nn.Sequential(nn.Linear(in_features=400, out_features=200),
                            nn.LeakyReLU(negative_slope=alpha))
    self.l3 = nn.Sequential(nn.Linear(in_features=200, out_features=200),
                            nn.LeakyReLU(negative_slope=alpha))
    # build output layer
    self.Qval = nn.Linear(in_features=200, out_features=24)

def forward(self, observation):
    if isinstance(observation, np.ndarray):
        observation = torch.from_numpy(observation).float()
    out1 = self.l1(observation)
    out2 = self.l2(out1)
    out3 = self.l3(out2)
    qval = self.Qval(out3)
    return qval

и код обратного распространения может быть, например:

self.optimizer = optim.Adam(self.q_net.parameters(), lr=1e-4)
self.optimizer.zero_grad()

state_batch=torch.rand([64,500])
act_batch=np.randi(0,24,[64,1]
act_batch_torch=torch.as_tensor(act_batch)
label_batch = torch.rand([64,500])
Q=self.q_net.forward(state_batch).gather(1, act_batch_torch) # q_net is an instance of the network above
loss = mse_loss(input=Q, target=label_batch.detach())
loss.backward()

self.optimizer.step()

Обратите внимание, что, поскольку логический вывод выполняется намного быстрее при использовании ЦП, я также делаю обратную передачу на ЦП. Я попытался передать сеть на графический процессор, а затем выполнить обратную передачу на графическом процессоре, но это оказалось медленнее.

Есть идеи, почему pyTorch медленнее? Как я могу улучшить скорость для этого типа неглубокой сети?