xref: /aosp_15_r20/external/pytorch/test/test_autoload.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# Owner(s): ["module: PrivateUse1"]
2
3import os
4
5from torch.testing._internal.common_utils import run_tests, TestCase
6
7
8class TestDeviceBackendAutoload(TestCase):
9    def test_autoload(self):
10        switch = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "0")
11
12        # After importing the extension, the value of this environment variable should be true
13        # See: test/cpp_extensions/torch_test_cpp_extension/__init__.py
14        is_imported = os.getenv("IS_CUSTOM_DEVICE_BACKEND_IMPORTED", "0")
15
16        # Both values should be equal
17        self.assertEqual(is_imported, switch)
18
19
20if __name__ == "__main__":
21    run_tests()
22