add the option to use a FORCE_CUDA to force cuda installation on docker

docker/Dockerfile

@@ -1,5 +1,6 @@
 ARG CUDA="9.0"
 ARG CUDNN="7"
+ARG FORCE_CUDA="1"
 
 FROM nvidia/cuda:${CUDA}-cudnn${CUDNN}-devel-ubuntu16.04
 
@@ -32,7 +33,6 @@ RUN conda install -y ipython
 RUN pip install ninja yacs cython matplotlib opencv-python
 
 # Install PyTorch 1.0 Nightly
-ARG CUDA
 RUN echo conda install pytorch-nightly cudatoolkit=${CUDA} -c pytorch \
  && conda clean -ya
 
@@ -47,8 +47,8 @@ RUN git clone https://github.com/cocodataset/cocoapi.git \
  && python setup.py build_ext install
 
 # install PyTorch Detection
-RUN git clone https://github.com/facebookresearch/maskrcnn-benchmark.git \
+RUN git clone https://github.com/bendidi/maskrcnn-benchmark.git \
  && cd maskrcnn-benchmark \
- && python setup.py build develop
+ && FORCE_CUDA=1 python setup.py build develop
 
 WORKDIR /maskrcnn-benchmark

setup.py

@@ -28,7 +28,8 @@ def get_extensions():
     extra_compile_args = {"cxx": []}
     define_macros = []
 
-    if torch.cuda.is_available() and CUDA_HOME is not None:
+    if (torch.cuda.is_available() and CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1":
+        print("\n\nINSTALLING WITH CUDA SUPPORT\n\n")
         extension = CUDAExtension
         sources += source_cuda
         define_macros += [("WITH_CUDA", None)]