有没有办法在Pytorch中并行使用torch.autograd.gradient？

Question

我正在尝试训练一些网络，其中损失不仅是输出的函数，而且是输出w.r.t的导数。输入。问题是，尽管可以使用Pytorh与模块并行执行批处理输出的计算，但我找不到并行执行导数的方法。这是我可以连续做的最好的事情：

import torch
x=torch.rand(300,1)
dydx=torch.zeros_like(x)
fc=torch.nn.Linear(1,1)
x.requires_grad=True
for ii in range(x.size(0)):
    xi=x[ii,0:] 
    yi=torch.tanh(fc(xi))    
    dydx[ii]=torch.autograd.grad(yi,xi,create_graph=True)[0]
dydxsum=(dydx**2).sum()
dydxsum.backward()

在上面的代码中，x被拆分以节省内存和时间。但是，当x的大小变大时，仍然需要并行化（在CUDA中）。如果必须通过修补Pytorch来实现，则将对从何处开始的提示表示赞赏。