More

Author: jrbourbeau

Dockerfile

@@ -10,6 +10,12 @@ WORKDIR /home
 COPY pytorch.yml pytorch.yml
 RUN mamba env update -n base --file pytorch.yml \
     && conda clean -afy \
+    && mamba uninstall -y pytorch torchvision \
+    && mamba install -y -n base -c conda-forge \
+        "cuda-version=${CUDA_VER%.*}.*" \
+        pytorch-gpu \
+        torchvision \
+    && conda clean -afy \
     && rm pytorch.yml