mightychaos / dqtorch Goto Github PK

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
from pkg_resources import DistributionNotFound, get_distribution, parse_version
import torch
_src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'dqtorch')

use_new_std = parse_version(torch.__version__) >= parse_version('2.0.0')

nvcc_flags = [
    '-O3', '-std=c++14' if not use_new_std else '-std=c++17',
    '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__',
]

if os.name == "posix":
    c_flags = ['-O3', '-std=c++14' if not use_new_std else '-std=c++17']
elif os.name == "nt":
...