Install pytorch-cpu for cpu image (kubeflow#536)

magdalenakuhn17 · Nov 10, 2019 · 3041486 · 3041486
1 parent 7f24834
commit 3041486
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/pytorch.Dockerfile b/python/pytorch.Dockerfile
@@ -17,12 +17,12 @@ RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-4.
      /opt/conda/bin/conda create -y --name pytorch-py37 python=3.7.3 numpy pyyaml scipy ipython mkl&& \
      /opt/conda/bin/conda clean -ya
 ENV PATH /opt/conda/envs/pytorch-py37/bin:$PATH
-RUN conda install --name pytorch-py37 pytorch torchvision -c soumith && /opt/conda/bin/conda clean -ya
+RUN conda install --name pytorch-py37 pytorch-cpu torchvision -c soumith && /opt/conda/bin/conda clean -ya
 
 WORKDIR /workspace
 RUN chmod -R a+w /workspace
 
 COPY . .
 RUN pip install --upgrade pip && pip install -e ./kfserving
 RUN pip install -e ./pytorchserver
-ENTRYPOINT ["python", "-m", "pytorchserver"]
+ENTRYPOINT ["python", "-m", "pytorchserver"]