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 11 12 13pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) 14sys.path.append(pytorch_test_dir) 15 16from dynamo.test_cudagraphs import TestAotCudagraphs # noqa: F401 17 18 19REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent 20 21sys.path.insert(0, str(REPO_ROOT)) 22from tools.stats.import_test_stats import get_disabled_tests 23 24 25# Make sure to remove REPO_ROOT after import is done 26sys.path.remove(str(REPO_ROOT)) 27 28if __name__ == "__main__": 29 if torch.cuda.is_available() and not IS_JETSON and not IS_WINDOWS: 30 get_disabled_tests(".") 31 32 torch.cuda.memory._set_allocator_settings("expandable_segments:True") 33 34 run_tests() 35