Pytorch-加速此方程式链

Question

    for idj in range(1, self.links.NUM_JOINTS):
        ipar = self.links.m_parent[idj]
        t_here = J_vec[idj, :] - J_vec[ipar, :]
        A_here = make_TransformMatrix(Rs[idj, :], j_here)
        outT[idj] = torch.matmul(outT[ipar], A_here)

我碰巧正在使用Pytorch对机器人链接进行优化。正如yu从上面的代码中可以看到的那样，我正在构建一种运动树，或将机器人各部分链接在一起的链。这很好用，我能够解决我的问题。

但是，此链有50多个关节，每个循环大约需要1毫秒，使其变为50毫秒，如果我迭代几次迭代，它会变得非常慢。减速的原因似乎是因为我必须“链接”梯度，因此必须通过该运动学链来方程式。

是否可以通过某种方式对此进行编码，因此它已经“预先计算”了该链，但是仅允许我更改输入到模型中的值，以便从中获得梯度？我不知道该如何处理。