Unverified Commit c9ba50a9 authored by Gustaf Ahdritz's avatar Gustaf Ahdritz Committed by GitHub
Browse files

Merge pull request #275 from zrqiao/zrqiao_cpuonly_setup

Resolving torch cuda availability for cpu-only installation
parents ca7ba44c bb11fb18
# Originally from Openfold https://github.com/aqlaboratory/openfold, modified.
# Copyright 2021 AlQuraishi Laboratory # Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited # Copyright 2021 DeepMind Technologies Limited
# #
...@@ -16,6 +17,7 @@ import os ...@@ -16,6 +17,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 +39,7 @@ extra_cuda_flags = [ ...@@ -37,7 +39,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