diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
index 9c653a737..a8c170d91 100644
--- a/.github/workflows/actions.yml
+++ b/.github/workflows/actions.yml
@@ -29,7 +29,7 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install -e ".[tests]" --progress-bar off --upgrade
+        pip install -e ".[tests]" --progress-bar off
     - name: Test with pytest
       run: |
         pytest --cov=autokeras --cov-report xml:coverage.xml
diff --git a/autokeras/engine/tuner.py b/autokeras/engine/tuner.py
index 59b3b4895..521318c0c 100644
--- a/autokeras/engine/tuner.py
+++ b/autokeras/engine/tuner.py
@@ -117,10 +117,10 @@ def get_output_layer(tensor):
             for layer in model.layers:
                 if isinstance(layer, tf.keras.layers.InputLayer):
                     continue
-                if not isinstance(layer, preprocessing.PreprocessingLayer):
-                    break
                 input_node = nest.flatten(layer.input)[0]
                 if input_node is tensor:
+                    if not isinstance(layer, preprocessing.PreprocessingLayer):
+                        break
                     return layer
             return None
 
diff --git a/setup.py b/setup.py
index 025052cb0..6dc7b127a 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,7 @@
     install_requires=[
         "packaging",
         "keras-tuner>=1.0.2",
-        "tensorflow>=2.3.0",
+        "tensorflow<=2.5.0,>=2.3.0",
         "scikit-learn",
         "pandas",
     ],