From 0e913b1c2efbf2901de3ef93e1d41a9c6882f3b4 Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Mon, 8 Aug 2022 04:16:26 -0700 Subject: [PATCH] PYTORCH3D_NO_EXTENSION Summary: Request in https://github.com/facebookresearch/pytorch3d/issues/1233 for option to disable CUDA build. Also option to disable binary build completely. This could be useful e.g. in the config tutorial where we need a small python-only part of pytorch3d. Reviewed By: kjchalup Differential Revision: D38458624 fbshipit-source-id: 421a0b1cc31306d7e322d3e743e30a7533d7f034 --- setup.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 3541bcbc..357073ab 100755 --- a/setup.py +++ b/setup.py @@ -8,6 +8,7 @@ import glob import os import runpy +import sys import warnings from typing import List, Optional @@ -35,6 +36,13 @@ def get_existing_ccbin(nvcc_args: List[str]) -> Optional[str]: def get_extensions(): + no_extension = os.getenv("PYTORCH3D_NO_EXTENSION", "0") == "1" + if no_extension: + msg = "SKIPPING EXTENSION BUILD. PYTORCH3D WILL NOT WORK!" + print(msg, file=sys.stderr) + warnings.warn(msg) + return [] + this_dir = os.path.dirname(os.path.abspath(__file__)) extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc") sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"), recursive=True) @@ -46,7 +54,10 @@ def get_extensions(): include_dirs = [extensions_dir] force_cuda = os.getenv("FORCE_CUDA", "0") == "1" - if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda: + force_no_cuda = os.getenv("PYTORCH3D_FORCE_NO_CUDA", "0") == "1" + if ( + not force_no_cuda and torch.cuda.is_available() and CUDA_HOME is not None + ) or force_cuda: extension = CUDAExtension sources += source_cuda define_macros += [("WITH_CUDA", None)]