diff --git a/setup.py b/setup.py index e2462ae13..261a620a3 100644 --- a/setup.py +++ b/setup.py @@ -64,36 +64,29 @@ def get_extensions(): use_cuda = torch.cuda.is_available() and CUDA_HOME is not None extension = CUDAExtension if use_cuda else CppExtension + extra_link_args = [] + extra_compile_args = { + "nvcc": [ + "-O3" if not debug_mode else "-O0", + "-t=0", + ] + } + if not IS_WINDOWS: - extra_link_args = [] - extra_compile_args = { - "cxx": [ - "-O3" if not debug_mode else "-O0", - "-fdiagnostics-color=always", - ], - "nvcc": [ - "-O3" if not debug_mode else "-O0", - "-t=0", - ] - } + extra_compile_args["cxx"] = [ + "-O3" if not debug_mode else "-O0", + "-fdiagnostics-color=always", + ] if debug_mode: extra_compile_args["cxx"].append("-g") extra_compile_args["nvcc"].append("-g") extra_link_args.extend(["-O0", "-g"]) - else: - extra_link_args = [] - extra_compile_args = { - "cxx": [ - "/O2" if not debug_mode else "/Od", - "/permissive-" - ], - "nvcc": [ - "-O3" if not debug_mode else "-O0", - "-t=0", - ] - } + extra_compile_args["cxx"] = [ + "/O2" if not debug_mode else "/Od", + "/permissive-" + ] if debug_mode: extra_compile_args["cxx"].append("/ZI") @@ -125,6 +118,7 @@ def get_extensions(): return ext_modules + setup( name="torchao", version=version+version_suffix, @@ -138,7 +132,7 @@ def get_extensions(): description="Package for applying ao techniques to GPU models", long_description=open("README.md").read(), long_description_content_type="text/markdown", - url="https://github.com/pytorch-labs/ao", + url="https://github.com/pytorch/ao", cmdclass={"build_ext": BuildExtension}, options={"bdist_wheel": { "py_limited_api": "cp39"