Fix NCCL_ASYNC_ERROR_HANDLING deprecation warning (#711)

* Fix NCCL_ASYNC_ERROR_HANDLING deprecation warning

It looks like the patch from pytorch/pytorch#114077 landed in torch 2.2.0.

Fixes #568.

* Update CHANGELOG.md

Author: simonbyrne

CHANGELOG.md

@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Update pytests to skip when the required dependencies are not present
 - Bug in data processing script in domino training example
+- Fixed NCCL_ASYNC_ERROR_HANDLING deprecation warning
 
 ### Security
 

modulus/distributed/manager.py

@@ -332,7 +332,11 @@ def initialize():
         addr = os.getenv("MASTER_ADDR", "localhost")
         port = os.getenv("MASTER_PORT", "12355")
         # https://pytorch.org/docs/master/notes/cuda.html#id5
-        os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
+        # was changed in version 2.2
+        if torch.__version__ < (2, 2):
+            os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
+        else:
+            os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "0"
         initialization_method = os.getenv("MODULUS_DISTRIBUTED_INITIALIZATION_METHOD")
         if initialization_method is None:
             try: