diff --git a/.github/workflows/full_tests.yml b/.github/workflows/full_tests.yml
index ad222b04f..7962d4a3a 100644
--- a/.github/workflows/full_tests.yml
+++ b/.github/workflows/full_tests.yml
@@ -85,7 +85,7 @@ jobs:
           python-version: ${{ env.PYTHON_VERSION }}
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot --no-extra wallx # TODO(Steven): Make flash-attn optional
+        run: uv sync --extra all # TODO(Steven): Make flash-attn optional
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv --maxfail=10
diff --git a/.github/workflows/unbound_deps_tests.yml b/.github/workflows/unbound_deps_tests.yml
index 95562d0dd..bf93dcda6 100644
--- a/.github/workflows/unbound_deps_tests.yml
+++ b/.github/workflows/unbound_deps_tests.yml
@@ -78,7 +78,7 @@ jobs:
           echo "Dependencies unbound:" && cat pyproject.toml
 
       - name: Install lerobot with all extras
-        run: uv sync --all-extras --no-extra groot --no-extra wallx # TODO(Steven): Make flash-attn optional
+        run: uv sync --extra all # TODO(Steven): Make flash-attn optional
 
       - name: Run pytest (all extras)
         run: uv run pytest tests -vv
diff --git a/pyproject.toml b/pyproject.toml
index 06b00e96c..ea84eec36 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -169,7 +169,7 @@ all = [
     "lerobot[kinematics]",
     "lerobot[intelrealsense]",
     # "lerobot[wallx]",
-    "lerobot[pi]",
+    # "lerobot[pi]", TODO(Pepijn): Update pi to transformers v5
     "lerobot[smolvla]",
     # "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
     "lerobot[xvla]",
@@ -406,6 +406,10 @@ conflicts = [
         { extra = "wallx" },
         { extra = "xvla" },
     ],
+    [
+        { extra = "wallx" },
+        { extra = "sarm" },
+    ],
     [
         { extra = "wallx" },
         { extra = "hilserl" },
@@ -418,4 +422,37 @@ conflicts = [
         { extra = "wallx" },
         { extra = "all" },
     ],
+    # pi uses custom branch which conflicts with transformers-dep
+    [
+        { extra = "pi" },
+        { extra = "transformers-dep" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "smolvla" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "groot" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "xvla" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "sarm" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "hilserl" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "libero" },
+    ],
+    [
+        { extra = "pi" },
+        { extra = "all" },
+    ],
 ]
diff --git a/src/lerobot/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py
index 0d9c77e00..e4970dcf1 100644
--- a/src/lerobot/policies/pi0/modeling_pi0.py
+++ b/src/lerobot/policies/pi0/modeling_pi0.py
@@ -93,10 +93,11 @@ def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedd
 
 
 def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    alpha_t = torch.as_tensor(alpha, dtype=torch.float32, device=device)
-    beta_t = torch.as_tensor(beta, dtype=torch.float32, device=device)
+    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
+    alpha_t = torch.tensor(alpha, dtype=torch.float32)
+    beta_t = torch.tensor(beta, dtype=torch.float32)
     dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,))
+    return dist.sample((bsize,)).to(device)
 
 
 def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)