Skip to content

GitLab

Commit 310f240c authored by Zhouran Qiao zhuoran.qiao@entos.ai's avatar Zhouran Qiao zhuoran.qiao@entos.ai
Browse files

Resolving torch cuda availablity for cpu-only installation

parent b2d6bff6
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
setup.py
View file @ 310f240c
...@@ -16,6 +16,7 @@ import os ...@@ -16,6 +16,7 @@ import os
from setuptools import setup, Extension, find_packages from setuptools import setup, Extension, find_packages
import subprocess import subprocess
import torch
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
from scripts.utils import get_nvidia_cc from scripts.utils import get_nvidia_cc
...@@ -37,7 +38,7 @@ extra_cuda_flags = [ ...@@ -37,7 +38,7 @@ extra_cuda_flags = [
] ]
def get_cuda_bare_metal_version(cuda_dir): def get_cuda_bare_metal_version(cuda_dir):
if cuda_dir==None: if cuda_dir==None or torch.version.cuda==None:
print("CUDA is not found, cpu version is installed") print("CUDA is not found, cpu version is installed")
return None, -1, 0 return None, -1, 0
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment