No need to ignore type for torch.__version__ as of PyTorch 1.6.0

HazyResearch · Jul 30, 2020 · 5909011 · 5909011
1 parent 1d6771b
commit 5909011
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 4 deletions.
diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml
@@ -34,8 +34,6 @@ jobs:
         sudo rm -rf /etc/ImageMagick-6/policy.xml
         sudo apt-get install -q -y poppler-utils
         pip install --upgrade pip setuptools
-        # Install PyTorch for Linux with no CUDA support
-        pip install torch==1.5.0+cpu torchvision==0.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
     - name: Print Version Info
       run: |
         neofetch

diff --git a/Dockerfile b/Dockerfile
@@ -28,7 +28,7 @@ ARG FONDUER_VERSION=
 
 # Install python packages
 RUN pip install \
-    torch==1.4.0+cpu -f https://download.pytorch.org/whl/torch_stable.html \
+    torch==1.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html \
     # https://github.com/HazyResearch/fonduer/issues/390
     "tensorboard<2.0.0,>=1.14.0" "scikit-learn<0.22.0,>=0.20.2" \
     fonduer${FONDUER_VERSION:+==${FONDUER_VERSION}}
diff --git a/src/fonduer/packaging/fonduer_model.py b/src/fonduer/packaging/fonduer_model.py
@@ -357,7 +357,7 @@ def _get_default_conda_env() -> Optional[Dict[str, Any]]:
 
     return _mlflow_conda_env(
         additional_conda_deps=[
-            "pytorch={}".format(torch.__version__),  # type: ignore
+            "pytorch={}".format(torch.__version__),
             "psycopg2",
             "pip",
         ],