diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 5d2f1e7b2..5ea97f8f3 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -60,6 +60,7 @@ jobs:
           pip install pytest hypothesis coverage coveralls
           pip install git+https://github.com/XENONnT/base_environment.git --force-reinstall
           pip install git+https://github.com/AxFoundation/strax.git --force-reinstall
+          pip install . --force-reinstall
 
       - name: Start MongoDB
         uses: supercharge/mongodb-github-action@1.11.0
@@ -101,6 +102,7 @@ jobs:
         run: |
           coverage run --source=straxen -m pytest --durations 0
           coverage report
+
       - name: Coveralls
         # Make the coverage report and upload
         env:
@@ -126,5 +128,6 @@ jobs:
           bash .github/scripts/create_pre_apply_function.sh $HOME
           coverage run --append --source=straxen -m pytest -v
           coveralls --service=github
+
       - name: goodbye
         run: echo 'tests done, bye bye'
diff --git a/straxen/config/protocols.py b/straxen/config/protocols.py
index 4d031d445..c7f815e07 100644
--- a/straxen/config/protocols.py
+++ b/straxen/config/protocols.py
@@ -97,7 +97,6 @@ def open_neural_net(model_path: str, custom_objects=None, **kwargs):
     """Open a tensorflow file and return a keras model."""
     # Nested import to reduce loading time of import straxen and it not
     # base requirement
-    import jax  # noqa: F401
     import tensorflow as tf
 
     if not os.path.exists(model_path):
@@ -273,7 +272,6 @@ def open_jax_model(model_path: str, **kwargs):
     """
     # Nested import to reduce loading time of import straxen as it's not
     # a base requirement
-    import jax  # noqa: F401
     from jax import export
 
     # Check if the model file exists