diff --git a/setup.py b/setup.py index 7e15e6266..46e9b53bc 100644 --- a/setup.py +++ b/setup.py @@ -93,6 +93,9 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int else: os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5" +if "TORCH_CUDA_ARCH_LIST" in os.environ: + print("\n\nBuilding Apex for archs {}\n\n".format(os.environ["TORCH_CUDA_ARCH_LIST"])) + print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) TORCH_MAJOR = int(torch.__version__.split(".")[0]) TORCH_MINOR = int(torch.__version__.split(".")[1])