deepfold.testing.common.with_nccl_blocking_wait¶

deepfold.testing.common.with_nccl_blocking_wait(func)[source]¶: Convenience decorator to set/unset TORCH_NCCL_BLOCKING_WAIT flag. Note that use of this decorator will override the setting of TORCH_NCCL_ASYNC_ERROR_HANDLING for the particular test. After the test, both TORCH_NCCL_BLOCKING_WAIT and TORCH_NCCL_ASYNC_ERROR_HANDLING will be restored to their original values.