Fix GPU docker builds?

Author: jrbourbeau

.github/workflows/docker.yml

@@ -2,8 +2,8 @@ name: Docker build
 
 on:
   push:
-    branches:
-      - main
+    # branches:
+    #   - main
   pull_request:
   workflow_dispatch:
 

Dockerfile

@@ -11,12 +11,11 @@ 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 pytorch -c nvidia -c conda-forge \
+    && mamba install -y -n base -c conda-forge \
         "cudatoolkit=${CUDA_VER%.*}.*" \
         "cuda-version=${CUDA_VER%.*}.*" \
-        pytorch \
-        torchvision \
-        "pytorch-cuda=${CUDA_VER%.*}.*" \
+        pytorch-gpu \
+        torchvision
     && conda clean -afy \
     && rm pytorch.yml