diff --git a/.github/workflows/doctests.yml b/.github/workflows/doctests.yml
index f8859e1c61b..cc95c90bcf0 100644
--- a/.github/workflows/doctests.yml
+++ b/.github/workflows/doctests.yml
@@ -34,7 +34,7 @@ jobs:
           nvidia-smi
 
       - name: Install transformers in edit mode
-        run: python3 -m pip install -e .
+        run: python3 -m pip install -e .[flax]
 
       - name: GPU visibility
         run: |