1# Owner(s): ["module: cuda"] 2# run time cuda tests, but with the allocator using expandable segments 3 4import os 5import pathlib 6import sys 7 8import torch 9from torch.testing._internal.common_cuda import IS_JETSON, IS_WINDOWS 10from torch.testing._internal.common_utils import run_tests, TEST_WITH_ASAN 11from torch.testing._internal.inductor_utils import HAS_CUDA 12 13 14pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) 15sys.path.append(pytorch_test_dir) 16 17if HAS_CUDA and not TEST_WITH_ASAN: 18 try: 19 from .test_cudagraph_trees import CudaGraphTreeTests 20 except ImportError: 21 from test_cudagraph_trees import CudaGraphTreeTests # noqa: F401 22 23REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent 24 25sys.path.insert(0, str(REPO_ROOT)) 26from tools.stats.import_test_stats import get_disabled_tests 27 28 29# Make sure to remove REPO_ROOT after import is done 30sys.path.remove(str(REPO_ROOT)) 31 32if __name__ == "__main__": 33 if ( 34 torch.cuda.is_available() 35 and not IS_JETSON 36 and not IS_WINDOWS 37 and HAS_CUDA 38 and not TEST_WITH_ASAN 39 ): 40 get_disabled_tests(".") 41 42 torch.cuda.memory._set_allocator_settings("expandable_segments:True") 43 44 run_tests() 45