add pytorch (#1551)

dmlc · Dec 1, 2020 · d9feba6 · d9feba6
1 parent 2429845
commit d9feba6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/tools/batch/docker/gluon_cv_job.sh b/tools/batch/docker/gluon_cv_job.sh
@@ -22,8 +22,10 @@ git fetch origin $SOURCE_REF:working
 git checkout working
 if [ $DEVICE == "cpu" ]; then
 	python3 -m pip install -U --quiet "mxnet==1.7.0.post1"
+    python3 -m pip install -U --quiet torch==1.6.0+cpu torchvision==0.7.0+cpu
 else
 	python3 -m pip install -U --quiet "mxnet-cu102==1.7.0"
+    python3 -m pip install -U --quiet torch==1.6.0 torchvision==0.7.0
 fi;
 
 python3 -m pip install --quiet -e .

diff --git a/tools/docker/dockerfile b/tools/docker/dockerfile
@@ -101,6 +101,9 @@ RUN pip3 install --no-cache --upgrade \
 # Install MXNet
 RUN pip3 install --no-cache --upgrade mxnet-cu102==1.7.0
 
+# Install PyTorch
+RUN pip3 install torch==1.6.0 torchvision==0.7.0
+
 RUN mkdir -p ${WORKDIR}/notebook
 RUN mkdir -p ${WORKDIR}/data
 RUN mkdir -p /.init