From 9f1d7df370495ade5ee46bdf19c9d431b4110e00 Mon Sep 17 00:00:00 2001 From: Francesco Innocenti Date: Fri, 28 Jun 2024 13:59:10 +0100 Subject: [PATCH] Set `energies_iters` to max compatible with `diffrax` --- jpc/_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jpc/_utils.py b/jpc/_utils.py index be9477a..93d0340 100644 --- a/jpc/_utils.py +++ b/jpc/_utils.py @@ -119,7 +119,7 @@ def loop_body(state): energies_iters = energies_iters.at[:, t].set(energies) return t + 1, energies_iters - energies_iters = zeros((len(model), 1000)) + energies_iters = zeros((len(model), 4096)) _, energies_iters = jax.lax.while_loop( lambda state: state[0] < t_max, loop_body,