jax

07278c5b - update gpu ram per test