search

CDCgov / cfa-viral-lineage-model

Apache License 2.0
9 stars 0 forks source link

CUDA warning #20

Closed swo closed 1 month ago

swo commented 1 month ago 
CUDA-enabled jaxlib is not installed

Does this mean we have the wrong version of jax in the lockfile?

afmagee42 commented 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.

swo commented 1 month ago

But my GFE laptop has a Nvidia T600 GPU! I want to use that power!

afmagee42 commented 1 month ago

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.)