mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
312652c325
Fixes #122468 - Load device extensions at the end of `torch/__init__.py` - Enabled by default, or you can disable it with `TORCH_DEVICE_BACKEND_AUTOLOAD=0` run test: ```python python test/run_test.py -i test_autoload_enable python test/run_test.py -i test_autoload_disable ``` doc: https://docs-preview.pytorch.org/pytorch/pytorch/127074/miscellaneous_environment_variables.html co-author: @jgong5 @bsochack @bkowalskiINTEL @jczaja @FFFrog @hipudding Co-authored-by: albanD <desmaison.alban@gmail.com> Co-authored-by: Jiong Gong <jiong.gong@intel.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/127074 Approved by: https://github.com/albanD, https://github.com/jgong5
22 lines
628 B
Python
22 lines
628 B
Python
# Owner(s): ["module: PrivateUse1"]
|
|
|
|
import os
|
|
|
|
from torch.testing._internal.common_utils import run_tests, TestCase
|
|
|
|
|
|
class TestDeviceBackendAutoload(TestCase):
|
|
def test_autoload(self):
|
|
switch = os.getenv("TORCH_DEVICE_BACKEND_AUTOLOAD", "0")
|
|
|
|
# After importing the extension, the value of this environment variable should be true
|
|
# See: test/cpp_extensions/torch_test_cpp_extension/__init__.py
|
|
is_imported = os.getenv("IS_CUSTOM_DEVICE_BACKEND_IMPORTED", "0")
|
|
|
|
# Both values should be equal
|
|
self.assertEqual(is_imported, switch)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|