saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[FIX] using torch.version.cuda/hip to ensure build ORTModule Torch C++ CUDA extension for docker build (#11675)

Browse source 
* [FIX] cpp ext

* Update orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py

Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com>

* [FIX] fix python format

Co-authored-by: Thiago Crepaldi <thiago.crepaldi@microsoft.com>
This commit is contained in:
PeixuanZuo 2022-06-07 07:51:26 +08:00 committed by GitHub
parent 981d45d8d5
commit 908e19dc16
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py
View file

@ -41,7 +41,7 @@ def _install_extension(ext_name, ext_path, cwd):
def build_torch_cpp_extensions():
"""Builds PyTorch CPP extensions and returns metadata."""
# Run this from within onnxruntime package folder
is_gpu_available = torch.cuda.is_available() and (
is_gpu_available = (torch.version.cuda is not None or torch.version.hip is not None) and (
ortmodule.ONNXRUNTIME_CUDA_VERSION is not None or ortmodule.ONNXRUNTIME_ROCM_VERSION is not None
)
os.chdir(ortmodule.ORTMODULE_TORCH_CPP_DIR)