use block_until_ready on benchmark output

valence-labs · Oct 1, 2024 · f3ad64b · f3ad64b
1 parent 48e9832
commit f3ad64b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/test_benchmark.py b/test/test_benchmark.py
@@ -68,6 +68,7 @@ def harness():
 def test_construct_basis(benchmark):
     def harness():
         mol = molecule("water")
-        basisset(mol, "6-31g")
+        basis = basisset(mol, "6-31g")
+        return jax.block_until_ready(basis)
 
     benchmark(harness)