Merge pull request #27 from instadeepai/fix/prioritised-sampling-with…

…-batch-of-1

fix: be explicit about squeeze dim in prioritised sampling to avoid flattening (1,1) arrays

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)
     )