Out of Memory error when RAM is available

[somearthling]

I'm running some large batches on a personal computer, with 16GB RAM and 4GB of VRAM, but I've run into a problem: When I run my code through the CUDA-enabled jaxlib, I run into an error where it says it couldn't allocate enough memory, and my VRAM seems to have topped out at 3.6 GB, but I'm only using about half my RAM. The code runs fine when I don't use CUDA and run it purely on my CPU. Concurrently to this, when running smaller batches, the first iteration takes much longer(3x) on the CUDA enabled version than the CPU version, while later iterations run faster on the GPU by about 2x. Is there any way to get the best of both worlds, and also avoid the memory error with larger batches?

[YouJiacheng]

OOM can refer to GPU out of memory.
You may manually transfer data between RAM and VRAM using device_put(to VRAM) and device_get(to RAM). I haven' tried it.

[somearthling]

In this code below:

probs = jax.vmap(circuit, in_axes=(0, None))(x, p)
jacobians = jax.vmap(jax.jacobian(circuit, argnums=1), in_axes=(0, None))(x, p)

would it help to device_get the array probs before the second line? If I do, can I do operations between probs and jacobians or will i have to use device_put probs back??

[YouJiacheng]

maybe you can replace jacobians = jax.vmap(jax.jacobian(circuit, argnums=1), in_axes=(0, None))(x, p) with

jacobians = jax.lax.map(lambda x: jax.jacobian(circuit, argnums=1)(x, p), x)