distributed computing (DDP) doesn't work for Hessian

[vganapati]

Possible resources:

• https://github.com/pytorch/functorch/issues/831
• https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html
• https://github.com/google/jax/issues/2937
• https://pytorch.org/tutorials/intermediate/jacobians_hessians.html
• https://github.com/pytorch/pytorch/issues/102197
• https://github.com/pytorch/pytorch/blob/main/functorch/examples/dp_cifar10/cifar10_transforms.py

[vganapati]

Fixed by implementing a custom DDP, see this issue.