Updated Dockerfile.

docker/Dockerfile

@@ -53,7 +53,8 @@ RUN apt-get update && apt-get install -y \
     libswresample-dev \
     less \
     groff \
-    mpich
+    mpich \
+    nvidia-cuda-dev
 
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
 
@@ -68,20 +69,9 @@ RUN curl https://rclone.org/install.sh | bash
 # Set timezone
 RUN ln -sf /usr/share/zoneinfo/US/Pacific /etc/localtime
 
-# Get Conda-ified Python.
-# RUN echo 'export PATH=/opt/conda/bin:$PATH' > /etc/profile.d/conda.sh && \
-#    wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
-#    sh ~/miniconda.sh -b -p /opt/conda && \
-#    rm ~/miniconda.sh
-
-# ENV PATH /opt/conda/bin:$PATH
-
 # Set CUDA_ROOT
 RUN export CUDA_HOME="/usr/local/cuda"
 
-# Install pytorch
-#RUN conda install pytorch torchvision cudatoolkit=11.1 -c pytorch -c conda-forge -y
-
 # Install zsh
 RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/v1.1.1/zsh-in-docker.sh)" -t robbyrussell -p git