Update Dockerfile

Dockerfile

@@ -1,9 +1,6 @@
 # Use an official PyTorch image with CUDA support as the base image
 FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
 
-# Set the CUDA_HOME environment variable
-ENV CUDA_HOME=/usr/local/cuda
-
 # Install Git and system libraries required for OpenGL without interactive prompts
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -20,6 +17,7 @@ RUN useradd -m -u 1000 user
 USER user
 
 ENV HOME=/home/user \
+    CUDA_HOME=/usr/local/cuda \
 	PATH=/home/user/.local/bin:$PATH \
     PYTHONPATH=$HOME/app \
 	PYTHONUNBUFFERED=1 \
@@ -65,8 +63,8 @@ RUN unzip -d $HOME/app/checkpoints/BFM $HOME/app/checkpoints/BFM.zip
 RUN find $HOME/app
 
 # Set the environment variable to specify the GPU device
-ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-ENV CUDA_VISIBLE_DEVICES=0
+#ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
+#ENV CUDA_VISIBLE_DEVICES=0
 
 # Run your app.py script
 CMD ["python", "app.py"]