From a5b4575e3c24d6f18a83691a5daf2430392228bd Mon Sep 17 00:00:00 2001
From: arition <aritionkb@gmail.com>
Date: Wed, 2 Dec 2020 00:01:21 -0800
Subject: [PATCH] Update PyTorch, torchvision and CUDA

---
 Dockerfile | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e78ba15..79c11b1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:10.2-cudnn7-runtime-centos8
+FROM nvidia/cuda:11.0-cudnn8-runtime-centos8
 
 RUN dnf -y install 'dnf-command(config-manager)' && \
     dnf -y config-manager --enable PowerTools && \
@@ -44,9 +44,11 @@ RUN dnf -y install nodejs && \
 
 RUN pip3 --no-cache-dir install --upgrade pip
 
+RUN pip3 --no-cache-dir install torch==1.7.0+cu110 torchvision==0.8.1+cu110 torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
+
 # tensorboard==2.2 has a breaking change that prevents jupyter_tensorboard==0.2.0 to function correctly
 # See more on: https://github.com/lspvic/jupyter_tensorboard/pull/63
-RUN pip3 --no-cache-dir install VapourSynth numpy tensorboard==2.1 jupyter jupyterlab torch torchvision efficientnet_pytorch jupyter-tensorboard pytorch-ignite ipywidgets Pillow Vizer matplotlib ipympl numba && \
+RUN pip3 --no-cache-dir install VapourSynth numpy tensorboard==2.1 jupyter jupyterlab efficientnet_pytorch jupyter-tensorboard pytorch-ignite ipywidgets Pillow Vizer matplotlib ipympl numba && \
     jupyter nbextension enable --py widgetsnbextension && \
     jupyter labextension install @jupyter-widgets/jupyterlab-manager && \
     jupyter labextension install jupyter-matplotlib && \
@@ -82,4 +84,4 @@ RUN mv /root/.jupyter /home/$NB_USER/ && \
 
 ENV HOME=/home/$NB_USER
 USER $NB_UID
-WORKDIR $HOME
\ No newline at end of file
+WORKDIR $HOME