From 2626599068b533d833fb2ad1052deba2e83b2a8d Mon Sep 17 00:00:00 2001 From: Daniel King Date: Wed, 24 Apr 2024 13:05:03 -0700 Subject: [PATCH] tdist --- scripts/train/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/train/train.py b/scripts/train/train.py index 2da61c46f6..89195b0dab 100644 --- a/scripts/train/train.py +++ b/scripts/train/train.py @@ -10,7 +10,7 @@ from typing import Any, Dict, List, Optional, Union import torch -import torch.distributed as tdist +import torch.distributed from composer import Trainer from composer.core.callback import Callback from composer.profiler import (JSONTraceHandler, Profiler, TraceHandler, @@ -132,7 +132,7 @@ def _initialize_gloo_and_nccl(dist_timeout: Union[int, float]): log.debug('Testing barrier with cpu...') dist.barrier() log.debug('Barrier test passed with cpu. Destroying process group...') - tdist.destroy_process_group() + torch.distributed.destroy_process_group() log.debug('Process group destroyed.') # Now, initialize with the correct device