Add back support for PYTORCH_TEST_WITH_MPS (#66)

Fix the TEST_WITH_MPS macro.
kulinseth · Feb 28, 2023 · 29026b3 · 29026b3
1 parent b7c2a65
commit 29026b3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py
@@ -13,7 +13,7 @@
 import torch.backends.mps
 from torch.testing._internal.common_utils import TestCase, TEST_WITH_ROCM, TEST_MKL, \
  skipCUDANonDefaultStreamIf, TEST_WITH_ASAN, TEST_WITH_UBSAN, TEST_WITH_TSAN, \
- IS_SANDCASTLE, IS_FBCODE, IS_REMOTE_GPU, IS_WINDOWS, \
+ IS_SANDCASTLE, IS_FBCODE, IS_REMOTE_GPU, IS_WINDOWS, TEST_WITH_MPS, \
  _TestParametrizer, compose_parametrize_fns, dtype_name, \
  NATIVE_DEVICES, skipIfTorchDynamo
 from torch.testing._internal.common_cuda import _get_torch_cuda_version, \
@@ -555,10 +555,8 @@ def get_device_type_test_bases():
  test_bases.append(CPUTestBase)
  if torch.cuda.is_available():
  test_bases.append(CUDATestBase)
- # Disable MPS testing in generic device testing temporarily while we're
- # ramping up support.
- # elif torch.backends.mps.is_available():
- # test_bases.append(MPSTestBase)
+ elif torch.backends.mps.is_available():
+ test_bases.append(MPSTestBase)
 
  return test_bases
 

diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
@@ -908,6 +908,7 @@ def _check_module_exists(name: str) -> bool:
 TEST_WITH_TSAN = os.getenv('PYTORCH_TEST_WITH_TSAN', '0') == '1'
 TEST_WITH_UBSAN = os.getenv('PYTORCH_TEST_WITH_UBSAN', '0') == '1'
 TEST_WITH_ROCM = os.getenv('PYTORCH_TEST_WITH_ROCM', '0') == '1'
+TEST_WITH_MPS = os.getenv('PYTORCH_TEST_WITH_MPS', '0') == '1'
 
 # Enables tests that are slow to run (disabled by default)
 TEST_WITH_SLOW = os.getenv('PYTORCH_TEST_WITH_SLOW', '0') == '1'