diff --git a/src/openpi/policies/droid_policy.py b/src/openpi/policies/droid_policy.py
index 666b7f76fe..55bdb419dd 100644
--- a/src/openpi/policies/droid_policy.py
+++ b/src/openpi/policies/droid_policy.py
@@ -78,4 +78,4 @@ def __call__(self, data: dict) -> dict:
 class DroidOutputs(transforms.DataTransformFn):
     def __call__(self, data: dict) -> dict:
         # Only return the first 8 dims.
-        return {"actions": np.asarray(data["actions"][:, :8])}
+        return {"actions": np.asarray(data["actions"][..., :8])}
diff --git a/src/openpi/policies/libero_policy.py b/src/openpi/policies/libero_policy.py
index 10611f61be..fe5aab0add 100644
--- a/src/openpi/policies/libero_policy.py
+++ b/src/openpi/policies/libero_policy.py
@@ -97,4 +97,4 @@ def __call__(self, data: dict) -> dict:
         # dimension, we need to now parse out the correct number of actions in the return dict.
         # For Libero, we only return the first 7 actions (since the rest is padding).
         # For your own dataset, replace `7` with the action dimension of your dataset.
-        return {"actions": np.asarray(data["actions"][:, :7])}
+        return {"actions": np.asarray(data["actions"][..., :7])}