From fd6d45f999b7e1f35ae19676ba23142fb024cc65 Mon Sep 17 00:00:00 2001 From: Brandon Franzke Date: Wed, 20 Sep 2017 22:28:44 -0700 Subject: [PATCH] add ENV variable to flag build with CUDA (instead of only "torch.cuda.is_available()") --- lib/layer_utils/roi_pooling/build.py | 5 ++++- lib/nms/build.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/lib/layer_utils/roi_pooling/build.py b/lib/layer_utils/roi_pooling/build.py index 620867b..f634645 100644 --- a/lib/layer_utils/roi_pooling/build.py +++ b/lib/layer_utils/roi_pooling/build.py @@ -8,7 +8,10 @@ defines = [] with_cuda = False -if torch.cuda.is_available(): +env_cuda = os.getenv('BUILD_WITH_CUDA', False) +print(env_cuda) + +if torch.cuda.is_available() or env_cuda: print('Including CUDA code.') sources += ['src/roi_pooling_cuda.c'] headers += ['src/roi_pooling_cuda.h'] diff --git a/lib/nms/build.py b/lib/nms/build.py index 4d9a96b..8076a93 100644 --- a/lib/nms/build.py +++ b/lib/nms/build.py @@ -8,7 +8,10 @@ defines = [] with_cuda = False -if torch.cuda.is_available(): +env_cuda = os.getenv('BUILD_WITH_CUDA', False) +print(env_cuda) + +if torch.cuda.is_available() or env_cuda: print('Including CUDA code.') sources += ['src/nms_cuda.c'] headers += ['src/nms_cuda.h']