From 0dc90dfa28ad7a016d4270d65c34dc45fd68509b Mon Sep 17 00:00:00 2001 From: Eitan Turok Date: Mon, 30 Sep 2024 17:28:32 +0000 Subject: [PATCH] clean up imports --- tests/tp/test_tp_strategies.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/tp/test_tp_strategies.py b/tests/tp/test_tp_strategies.py index 5ce63e78e4..77ec61de26 100644 --- a/tests/tp/test_tp_strategies.py +++ b/tests/tp/test_tp_strategies.py @@ -2,12 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 import os +import shutil from pathlib import Path -from shutil import rmtree from tempfile import TemporaryDirectory import pytest -from composer.utils import dist from omegaconf import OmegaConf as om from torch.distributed._tensor import Replicate, Shard from torch.distributed.tensor.parallel import ( @@ -122,7 +121,7 @@ def test_tp_train(tp_strategy: str): finally: # always remove data directory if os.path.isdir(data_dir): - rmtree(data_dir) + shutil.rmtree(data_dir) @pytest.mark.gpu