Computing Jacobian Vector Product (jvp) using nn.DataParallel

autograd
4

It seems the static jvp method is missing in an autograd.Function and you would need to implement it as described here.