Commit 6c2d04c2 authored by rusty1s's avatar rusty1s
Browse files

add FORCE_CUDA flag

parent 89c21e32
import os
import sys
import glob
import os
import os.path as osp
from setuptools import setup, find_packages
import sys
import torch
from setuptools import find_packages, setup
from torch.__config__ import parallel_info
from torch.utils.cpp_extension import BuildExtension
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension, CppExtension,
CUDAExtension)
WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
if os.getenv('FORCE_CUDA', '0') == '1':
WITH_CUDA = True
if os.getenv('FORCE_ONLY_CPU', '0') == '1':
WITH_CUDA = False
def get_extensions():
......
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