为什么这个tensorflow训练需要这么长时间？

我正在通过Deep Reinforcement Learning in Action一书学习 DRL 。在第 3 章中，他们展示了简单游戏 Gridworld（此处的说明，在规则部分）以及PyTorch 中的相应代码。

我已经对代码进行了试验，用 89% 的胜利（在训练后赢得 100 场比赛中的 89 场）训练网络只需不到 3 分钟。

作为练习，我已将代码迁移到tensorflow。所有代码都在这里。

问题是，使用我的 tensorflow 端口，以 84% 的胜率训练网络需要将近 2 个小时。两个版本都使用唯一的 CPU 进行训练（我没有 GPU）

训练损失数字似乎是正确的，也是获胜率（我们必须考虑到游戏是随机的，可能有不可能的状态）。问题是整个过程的性能。

我正在做一些非常错误的事情，但是什么？

主要区别在于训练循环，在火炬中是这样的：

        loss_fn = torch.nn.MSELoss()
        learning_rate = 1e-3
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
        ....
        Q1 = model(state1_batch) 
        with torch.no_grad():
            Q2 = model2(state2_batch) #B
        
        Y = reward_batch + gamma * ((1-done_batch) * torch.max(Q2,dim=1)[0])
        X = Q1.gather(dim=1,index=action_batch.long().unsqueeze(dim=1)).squeeze()
        loss = loss_fn(X, Y.detach())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在 tensorflow 版本中：

        loss_fn = tf.keras.losses.MSE
        learning_rate = 1e-3
        optimizer = tf.keras.optimizers.Adam(learning_rate)
        ...
        Q2 = model2(state2_batch) #B
        with tf.GradientTape() as tape:
            Q1 = model(state1_batch)
            Y = reward_batch + gamma * ((1-done_batch) * tf.math.reduce_max(Q2, axis=1))
            X = [Q1[i][action_batch[i]] for i in range(len(action_batch))]
            loss = loss_fn(X, Y)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

为什么培训需要这么长时间？