diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 6aa04b1..1d3c69b 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -21,6 +21,6 @@ jobs:
           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
       - name: Install build dependencies
         run: |
-          uv pip install build
+          uv pip install build ".[torch,flax]"
       - name: Build package
         run: python -m build
\ No newline at end of file
diff --git a/.github/workflows/manual_publish.yml b/.github/workflows/manual_publish.yml
index 473a742..99f4bbd 100644
--- a/.github/workflows/manual_publish.yml
+++ b/.github/workflows/manual_publish.yml
@@ -16,9 +16,9 @@ jobs:
         run: |
           curl -LsSf https://astral.sh/uv/install.sh | sh
           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
-      - name: Install Poetry
+      - name: Install Poetry and dependencies
         run: |
-          uv pip install poetry
+          uv pip install poetry ".[torch,flax]"
       - name: Build and Publish to PyPI
         run: |
           poetry config pypi-token.pypi ${{ secrets.PYPI_API_TOKEN }}