Enable support for complex gradient reduction in distributed cases.

NVIDIA / TorchFort

An Online Deep Learning Interface for HPC programs on NVIDIA GPUs

https://nvidia.github.io/TorchFort/

Other

154 stars 19 forks source link

Enable support for complex gradient reduction in distributed cases. #2

Closed romerojosh closed 1 year ago

romerojosh commented 1 year ago

This PR enables the distributed backend in TorchFort to handle complex valued gradients. Complex gradients can arise in models using Fourier neural operators or other model architectures that use complex weights.