Add CUDA_VISIBLE_DEVICES check (#1916)

Author: atalman

.github/scripts/validate_binaries.sh

@@ -69,6 +69,11 @@ else
         export PATH=${OLD_PATH}
     fi
 
+    # Use case CUDA_VISIBLE_DEVICES: https://github.com/pytorch/pytorch/issues/128819
+    if [[ ${MATRIX_GPU_ARCH_TYPE} == 'cuda' ]]; then
+        python3 -c "import torch;import os;print(torch.cuda.device_count(), torch.__version__);os.environ['CUDA_VISIBLE_DEVICES']='0';print(torch.empty(2, device='cuda'))"
+    fi
+
     # this is optional step
     if [[ ${TARGET_OS} != linux*  ]]; then
         conda deactivate

test/smoke_test/smoke_test.py

@@ -181,13 +181,6 @@ def smoke_test_cuda(package: str, runtime_error_check: str, torch_compile_check:
         if (sys.platform in ["linux", "linux2"]):
             print(f"torch nccl version: {torch.cuda.nccl.version()}")
 
-
-         # test settind CUDA_VISIBLE_DEVICES
-        print(f"torch.cuda.device_count: {torch.cuda.device_count()}")
-        os.environ['CUDA_VISIBLE_DEVICES']='0'
-        print(torch.empty(2, device='cuda'))
-        del os.environ['CUDA_VISIBLE_DEVICES']
-
         if runtime_error_check == "enabled":
             test_cuda_runtime_errors_captured()