Checkpoint commit was successful to /home/aklein4/outer-loop-jax/local_data/runs/baseline-fixed/checkpoints/000000000950/state