Update Dockerfile

Dockerfile

@@ -51,8 +51,8 @@ ENV PYTHONUNBUFFERED=1
 
 # Set the environment variable to specify the GPU device
 ENV CUDA_HOME=/bin/nvcc
-ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-ENV CUDA_VISIBLE_DEVICES=0
+# ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
+# ENV CUDA_VISIBLE_DEVICES=0
 
 # ENV PATH=${CUDA_HOME}/bin:/home/${USER_NAME}/.local/bin:/usr/bin:${PATH}
 # ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:/usr/local/lib/python3.8/site-packages/open3d:/usr/lib:/usr/lib64:${LD_LIBRARY_PATH}
@@ -73,13 +73,14 @@ RUN pip install --upgrade pip ninja
 RUN pip install setuptools==69.5.1
 RUN pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
 
-RUN python -c "import torch; print(torch.version.cuda)"
+RUN python -c "import torch, os; print(torch.version.cuda); print(os.environ.get('CUDA_PATH')); print(os.environ.get('CUDA_HOME'))"
 COPY requirements.txt /tmp
 RUN cd /tmp && pip install -r requirements.txt
 
 RUN pip install https://download.is.tue.mpg.de/icon/HF/kaolin-0.11.0-cp38-cp38-linux_x86_64.whl
 RUN pip install https://download.is.tue.mpg.de/icon/HF/pytorch3d-0.7.0-cp38-cp38-linux_x86_64.whl
 RUN pip install pyembree
+RUN pip install git+https://github.com/YuliangXiu/neural_voxelization_layer.git
 
 ENV TRANSFORMERS_CACHE=/tmp
 ENV MPLCONFIGDIR=/tmp