pytorch

1
# Owner(s): ["module: PrivateUse1"]
2

3
import os
4

5
from torch.testing._internal.common_utils import run_tests, TestCase
6

7

8
class 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

20
if __name__ == "__main__":
21
    run_tests()
22