Fix for Cuda 10.1 /Cuda 10.2 related torch package installation issue in Docker.

docker/Dockerfile

@@ -19,7 +19,7 @@
 ARG BASE_IMAGE=ubuntu:18.04
 
 FROM ${BASE_IMAGE} AS compile-image
-
+ARG BASE_IMAGE=ubuntu:18.04
 ENV PYTHONUNBUFFERED TRUE
 
 RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
@@ -41,13 +41,17 @@ RUN python3 -m venv /home/venv
 
 ENV PATH="/home/venv/bin:$PATH"
 
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1
-RUN update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 \
+    && update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1
 
-# This is only useful for cuda env 
+# This is only useful for cuda env
 RUN export USE_CUDA=1
 
-RUN pip install --no-cache-dir torch torchvision torchtext torchserve torch-model-archiver
+RUN if [ "$BASE_IMAGE" != "${BASE_IMAGE#*cuda:10.1}" ]; then \
+    pip install --no-cache-dir torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html; \
+    else \
+    pip install --no-cache-dir torch torchvision torchtext torchserve torch-model-archiver; \
+    fi
 
 # Final image for production
 FROM ${BASE_IMAGE} AS runtime-image

docker/README.md

@@ -29,7 +29,14 @@ For creating CPU based image :
 DOCKER_BUILDKIT=1 docker build --file Dockerfile -t torchserve:latest .
 ```
 
-For creating GPU based image :
+For creating GPU based image with Cuda 10.2 :
+
+```bash
+DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE=nvidia/cuda:10.2-cudnn7-runtime-ubuntu18.04 -t torchserve:latest .
+```
+
+For creating GPU based image with Cuda 10.1 :
+
 ```bash
 DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE=nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04 -t torchserve:latest .
 ```