mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
21 lines
626 B
Python
21 lines
626 B
Python
# Owner(s): ["module: ci"]
|
|
|
|
from torch.testing._internal.common_utils import TestCase, run_tests
|
|
|
|
|
|
# these tests could eventually be changed to fail if the import/init
|
|
# time is greater than a certain threshold, but for now we just use them
|
|
# as a way to track the duration of `import torch`.
|
|
class TestImportTime(TestCase):
|
|
def test_time_import_torch(self):
|
|
TestCase.runWithPytorchAPIUsageStderr("import torch")
|
|
|
|
def test_time_cuda_device_count(self):
|
|
TestCase.runWithPytorchAPIUsageStderr(
|
|
"import torch; torch.cuda.device_count()",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|