make `diff` more robust
@@ -17,5 +17,6 @@ def diff(x, t, order=1):
ones = torch.ones_like(x)
der, = autograd.grad(x, t, create_graph=True, grad_outputs=ones)
for i in range(1, order):
+ ones = torch.ones_like(der)
der, = autograd.grad(der, t, create_graph=True, grad_outputs=ones)
return der
