Fix vectorized axis for multi seed config

tum-pbs · Nov 5, 2024 · 912382d · 912382d
1 parent 5045acd
commit 912382d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/apebench/_base_scenario.py b/apebench/_base_scenario.py
@@ -832,7 +832,7 @@ def perform_tests_on_rollout(
                     jax.vmap(
                         func,
                         # Vectorize over time steps
-                        in_axes=(1, 2),
+                        in_axes=(1, 1),
                     ),
                     # Vectorize over seeds (but broadcast the reference)
                     in_axes=(0, None),