update docker to use cuda image

Dockerfile

@@ -1,4 +1,4 @@
-from ubuntu:20.04
+from nvidia/cuda:11.3.1-runtime-ubuntu20.04
 
 RUN apt update && \
     apt install -y bash \
@@ -11,7 +11,7 @@ RUN apt update && \
     rm -rf /var/lib/apt/lists
 
 RUN python3 -m pip install --no-cache-dir --upgrade pip && \
-    python3 -m pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \
+    python3 -m pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu113 \
     torch \
     torchvision \
     torchaudio