Closed swo closed 1 month ago
If that's from this:
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
I think JAX is just trying to be helpful and thus I don't think we care unless we want to run it on a GPU. And we don't.
But my GFE laptop has a Nvidia T600 GPU! I want to use that power!
A GPU is not a guarantee of a speed increase for any arbitrary model. It's easy for computational bottlenecks to form that wipe out any speed gains, and that's in cases where matrix-matrix multiplication is actually the rate-limiting step. (And where the code for interfacing with the GPU is smart and efficient in the first place.)
Does this mean we have the wrong version of jax in the lockfile?