Set cuda version to be None instead of an empty string (#9435)

This commit is contained in:
baijumeswani 2021-10-19 08:10:52 -07:00 committed by GitHub
parent e22920d954
commit 757bc66720
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 4 deletions

View file

@ -30,8 +30,8 @@ ORTMODULE_FALLBACK_POLICY = _FallbackPolicy.FALLBACK_UNSUPPORTED_DEVICE |\
ORTMODULE_FALLBACK_RETRY = False
ORTMODULE_IS_DETERMINISTIC = torch.are_deterministic_algorithms_enabled()
ONNXRUNTIME_CUDA_VERSION = ort_info.cuda_version if hasattr(ort_info, 'cuda_version') else ''
ONNXRUNTIME_ROCM_VERSION = ort_info.rocm_version if hasattr(ort_info, 'rocm_version') else ''
ONNXRUNTIME_CUDA_VERSION = ort_info.cuda_version if hasattr(ort_info, 'cuda_version') else None
ONNXRUNTIME_ROCM_VERSION = ort_info.rocm_version if hasattr(ort_info, 'rocm_version') else None
# Verify minimum PyTorch version is installed before proceding to ONNX Runtime initialization
try:

View file

@ -48,8 +48,8 @@ def build_torch_cpp_extensions():
os.chdir(ORTMODULE_TORCH_CPP_DIR)
# Extensions might leverage CUDA/ROCM versions internally
os.environ["ONNXRUNTIME_CUDA_VERSION"] = ONNXRUNTIME_CUDA_VERSION
os.environ["ONNXRUNTIME_ROCM_VERSION"] = ONNXRUNTIME_ROCM_VERSION
os.environ["ONNXRUNTIME_CUDA_VERSION"] = ONNXRUNTIME_CUDA_VERSION if not ONNXRUNTIME_CUDA_VERSION is None else ''
os.environ["ONNXRUNTIME_ROCM_VERSION"] = ONNXRUNTIME_ROCM_VERSION if not ONNXRUNTIME_ROCM_VERSION is None else ''
############################################################################
# Pytorch CPP Extensions that DO require CUDA/ROCM