diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml
index dc828c87..a6c915f5 100644
--- a/.github/workflows/build_docs.yml
+++ b/.github/workflows/build_docs.yml
@@ -1,39 +1,77 @@
-name: deploy
+name: docs
 
 on:
   push:
-    branches: # branch to trigger deployment
+    branches:
       - main
+    paths:
+      - "docs/**"
+      - "napari_cellseg3d/**"
+      - "pyproject.toml"
+      - ".github/workflows/docs.yml"
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "docs/**"
+      - "napari_cellseg3d/**"
+      - "pyproject.toml"
+      - ".github/workflows/docs.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: pages-${{ github.ref }}
+  cancel-in-progress: true
 
-# This job installs dependencies, build the book, and pushes it to `gh-pages`
 jobs:
-  build-and-deploy-book:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ubuntu-latest]
-        python-version: [3.9]
+  build:
+    name: build docs
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - name: Install docs dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e ".[docs]"
+
+      - name: Build Jupyter Book
+        run: |
+          jupyter-book build docs/
+
+      - name: Upload Pages artifact
+        if: github.event_name != 'pull_request'
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: docs/_build/html
+
+  deploy:
+    name: deploy docs
+    needs: build
+    runs-on: ubuntu-latest
+
+    if: github.event_name != 'pull_request' && github.ref == 'refs/heads/main'
+
+    permissions:
+      pages: write
+      id-token: write
+
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+
     steps:
-    - uses: actions/checkout@v2
-
-    # Install dependencies
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        pip install jupyter-book
-        pip install -e .
-
-    # Build the book
-    - name: Build the book
-      run: |
-        jupyter-book build docs/
-
-    # Deploy the book's HTML to gh-pages branch
-    - name: GitHub Pages action
-      uses: peaceiris/actions-gh-pages@v3.6.1
-      with:
-        github_token: ${{ secrets.GITHUB_TOKEN }}
-        publish_dir: docs/_build/html
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4
diff --git a/napari_cellseg3d/code_models/models/model_SwinUNetR.py b/napari_cellseg3d/code_models/models/model_SwinUNetR.py
index f90ef1df..8c59e93a 100644
--- a/napari_cellseg3d/code_models/models/model_SwinUNetR.py
+++ b/napari_cellseg3d/code_models/models/model_SwinUNetR.py
@@ -53,9 +53,15 @@ def __init__(
         try:
             parent_init(**init_kwargs)
         except TypeError as e:
-            logger.warning(f"Caught TypeError: {e}")
-            init_kwargs["in_channels"] = 1
-            parent_init(**init_kwargs)
+            if "img_size" in init_kwargs:
+                logger.warning(
+                    "Retrying SwinUNETR initialization without img_size due to "
+                    f"MONAI API compatibility issue: {e}"
+                )
+                init_kwargs.pop("img_size", None)
+                parent_init(**init_kwargs)
+            else:
+                raise
 
     # def forward(self, x_in):
     #     y = super().forward(x_in)
diff --git a/pyproject.toml b/pyproject.toml
index d2f79565..9d392834 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -148,7 +148,7 @@ dev = [
     "twine",
 ]
 docs = [
-    "jupyter-book",
+    "jupyter-book<2",
 ]
 test = [
     "pytest",