Oh, so just to make sure we are on the same page, the original OOM was during loading or training?
The archive I sent is precompiled using CUDA 12.1, rather than 12.0. According to Jax README it should still be compatible with CUDA 12.0, in which case we would be fine precompiling with 12.1. I will need to do some further tests.
