423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-11-04 18:02:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

PYTORCH3D_NO_EXTENSION

Browse Source 
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
This commit is contained in:
Jeremy Reizenstein 2022-08-08 04:16:26 -07:00 committed by Facebook GitHub Bot
parent bd93e9ce21
commit 0e913b1c2e
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
setup.py
View File

@ -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)]