1# Owner(s): ["module: ci"] 2 3from torch.testing._internal.common_utils import TestCase, run_tests 4 5 6# these tests could eventually be changed to fail if the import/init 7# time is greater than a certain threshold, but for now we just use them 8# as a way to track the duration of `import torch`. 9class TestImportTime(TestCase): 10 def test_time_import_torch(self): 11 TestCase.runWithPytorchAPIUsageStderr("import torch") 12 13 def test_time_cuda_device_count(self): 14 TestCase.runWithPytorchAPIUsageStderr( 15 "import torch; torch.cuda.device_count()", 16 ) 17 18 19if __name__ == "__main__": 20 run_tests() 21