diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index eb1aa6a2..77c1ee94 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -18,13 +18,13 @@ jobs:
           python-version: 3.7
       - name: Install project dependencies
         run: |
-          pip install tensorflow==2.1.0
+          pip install tensorflow-cpu==2.2.0
           pip install -e .[test]
       - name: Run Flake8
         run: flake8
       - name: Black code style
         run: black . --check --target-version py36 --exclude 'build/|buck-out/|dzzist/|_build/|\.git/|\.hg/|\.mypy_cache/|\.tox/|\.venv/|larq/snapshots/'
       - name: Check import order with isort
-        run: isort --check-only --diff
+        run: isort . --check --diff
       - name: Type check with PyType
-        run: pytype --jobs 2
+        run: pytype --jobs auto
diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index 272ba91f..cf8ca54e 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -11,8 +11,15 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.6, 3.7]
-        tf-version: [1.15.2, 2.0.1, 2.1.0, 2.2.0rc3]
+        tf-version: [1.15.3, 2.0.2, 2.1.1, 2.2.0]
+        python-version: [3.7]
+        include:
+          - tf-version: 2.2.0
+            python-version: 3.6
+          - tf-version: 2.2.0
+            python-version: 3.8
+          - tf-version: 2.2.0
+            python-version: 3.7
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python
@@ -21,13 +28,24 @@ jobs:
           python-version: ${{matrix.python-version}}
       - name: Install dependencies
         run: |
-          pip install tensorflow==${{matrix.tf-version}}
+          pip install tensorflow-cpu==${{matrix.tf-version}} || pip install tensorflow==${{matrix.tf-version}}
           pip install -e .[test]
       - name: Test data and models with pytest
         run: pytest . -n auto --ignore=tests/train_test.py --cov=larq_zoo --cov-report=xml --cov-config=.coveragerc
+      - name: Upload coverage to Codecov
+        run: bash <(curl -s https://codecov.io/bash) -f ./coverage.xml -F unittests
+
+  train-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2.0.1
+        with:
+          python-version: 3.7
+      - name: Install dependencies
+        run: pip install tensorflow-cpu==2.2.0 -e .[test]
       - name: Test training with pytest
-        if: matrix.python-version == '3.7' && matrix.tf-version == '2.1.0'
         run: pytest tests/train_test.py -n auto --cov=larq_zoo --cov-report=xml --cov-config=.coveragerc --cov-append
       - name: Upload coverage to Codecov
-        if: matrix.python-version == '3.7' && matrix.tf-version == '2.1.0'
-        run: curl -s https://codecov.io/bash | bash -s -- -f ./coverage.xml -F unittests
+        run: bash <(curl -s https://codecov.io/bash) -f ./coverage.xml -F unittests
diff --git a/larq_zoo/__init__.py b/larq_zoo/__init__.py
index d54d89df..c9bc79c8 100644
--- a/larq_zoo/__init__.py
+++ b/larq_zoo/__init__.py
@@ -1,5 +1,4 @@
-import larq_zoo.literature as literature
-import larq_zoo.sota as sota
+from larq_zoo import literature, sota
 from larq_zoo.core.utils import decode_predictions
 from larq_zoo.training import datasets
 from larq_zoo.training.data import preprocess_input
diff --git a/larq_zoo/core/utils.py b/larq_zoo/core/utils.py
index 39d35647..dc182c68 100644
--- a/larq_zoo/core/utils.py
+++ b/larq_zoo/core/utils.py
@@ -11,6 +11,7 @@
 try:
     from tensorflow.python.keras.applications.imagenet_utils import obtain_input_shape
 except ImportError:
+    # type: ignore
     from keras_applications.imagenet_utils import (
         _obtain_input_shape as obtain_input_shape,
     )
diff --git a/setup.cfg b/setup.cfg
index 49cfde28..21812883 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -7,17 +7,16 @@ exclude = build,dist,env,.pytype
 
 [isort]
 
-# Settings for compatibility with black.
+profile = black
 line_length = 88
-use_parentheses = True
-multi_line_output = 3
-include_trailing_comma = True
 # Don't misclassify larq as a first-party import.
 known_third_party = larq
 skip =
     build
     dist
-    env
+    venv
+    .venv
+    .git
     .pytype
 
 
diff --git a/setup.py b/setup.py
index 6f36988e..163ac383 100644
--- a/setup.py
+++ b/setup.py
@@ -24,14 +24,13 @@ def readme():
         "test": [
             "black==19.10b0",
             "flake8>=3.7.9,<3.9.0",
-            "isort~=4.3.21",
-            "pytype>=2019.10.17,<2020.7.0",
+            "isort==5.0.1",
+            "pytype==2020.6.26",
             "pytest>=4.3.1",
             "pytest-cov>=2.6.1",
             "pytest-mock>=3.1.1",
             "pytest-xdist==1.32.0",
             "Pillow==7.2.0",
-            "scipy==1.5.0",
             "tensorflow_datasets>=3.1.0",
         ],
     },