fix: be explicit about squeeze dim to avoid flattening (1,1) arrays.

instadeepai · Jun 14, 2024 · 2830156 · 2830156
1 parent b556e15
commit 2830156
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flashbax/buffers/sum_tree.py b/flashbax/buffers/sum_tree.py
@@ -208,7 +208,7 @@ def stratified_sample(
     )
 
     return jax.vmap(sample, in_axes=(None, None, 0))(
-        state, None, query_values.squeeze()
+        state, None, query_values.squeeze(axis=-1)
     )