frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Jetson Update for CI Redo (#94549)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94549
Approved by: https://github.com/ezyang, https://github.com/malfet
This commit is contained in:
puririshi98
2023-02-21 17:13:38 +00:00
committed by PyTorch MergeBot
parent c6d8d10b3e
commit 8aa34602f7
17 changed files with 94 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/test_cuda_nvml_based_avail.py
View File

@ -13,7 +13,7 @@ with patch.dict(os.environ, {"PYTORCH_NVML_BASED_CUDA_CHECK": "1"}):
# Before executing the desired tests, we need to disable CUDA initialization and fork_handler additions that would
# otherwise be triggered by the `torch.testing._internal.common_utils` module import
from torch.testing._internal.common_utils import (parametrize, instantiate_parametrized_tests, run_tests, TestCase,
IS_WINDOWS)
IS_WINDOWS, IS_JETSON)
# NOTE: Because `remove_device_and_dtype_suffixes` initializes CUDA context (triggered via the import of
# `torch.testing._internal.common_device_type` which imports `torch.testing._internal.common_cuda`) we need
# to bypass that method here which should be irrelevant to the parameterized tests in this module.
@ -48,6 +48,8 @@ class TestExtendedCUDAIsAvail(TestCase):
@parametrize("nvml_avail", [True, False])
@parametrize("avoid_init", ['1', '0', None])
def test_cuda_is_available(self, avoid_init, nvml_avail):
if IS_JETSON and nvml_avail and avoid_init == '1':
self.skipTest('Not working for Jetson')
patch_env = {"PYTORCH_NVML_BASED_CUDA_CHECK": avoid_init} if avoid_init else {}
with patch.dict(os.environ, **patch_env):
if nvml_avail: