Installing tochvision from official pip wheel (#127)

Author: Roshrini

docker/1.2.0/py2/Dockerfile.gpu

@@ -75,7 +75,6 @@ RUN conda install -c pytorch magma-cuda100==2.5.1 && \
                      requests==2.22.0 && \
     conda clean -ya
 
-ARG PYTORCH_VERSION=1.2.0
 ARG TORCHVISION_VERSION=0.4.0
 WORKDIR /opt/pytorch
 
@@ -84,7 +83,7 @@ ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
 ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
 
 RUN pip install --no-cache-dir https://pytorch-deep-learning-container.s3.amazonaws.com/torch-1.2.0-cp27-cp27mu-linux_x86_64.whl
-RUN pip install --no-cache-dir https://pytorch-deep-learning-container.s3.amazonaws.com/torchvision-0.4.0-cp27-cp27mu-linux_x86_64.whl
+RUN pip install torchvision==$TORCHVISION_VERSION
 
 WORKDIR /root
 

docker/1.2.0/py3/Dockerfile.gpu

@@ -75,8 +75,6 @@ RUN conda install -c pytorch magma-cuda100==2.5.1 && \
                      requests==2.22.0 && \
     conda clean -ya
 
-
-ARG PYTORCH_VERSION=1.2.0
 ARG TORCHVISION_VERSION=0.4.0
 WORKDIR /opt/pytorch
 
@@ -85,7 +83,7 @@ ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
 ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
 
 RUN pip install --no-cache-dir https://pytorch-deep-learning-container.s3.amazonaws.com/torch-1.2.0-cp36-cp36m-linux_x86_64.whl
-RUN pip install --no-cache-dir https://pytorch-deep-learning-container.s3.amazonaws.com/torchvision-0.4.0-cp36-cp36m-linux_x86_64.whl
+RUN pip install torchvision==$TORCHVISION_VERSION
 
 WORKDIR /root