4 ms·
It's useful to be able to run models/code on the CPU or split between CPU and GPU, especially for models that cannot fit into the GPU's VRAM, if you are running
by rhdunn 5mo ago
It's useful to be able to run models/code on the CPU or split between CPU and GPU, especially for models that cannot fit into the GPU's VRAM, if you are running multiple models, or if you have training data that you need to move between CPU and GPU.
- earthnail 5mo agoAll of that would be possible with the changes parent poster proposes. See the gotchas section in JAX, which is exactly these limitations: https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_J...