how to use gpu during training

[ZepSbosnia]

hi I'd like to know if it is possible to run the train on the GPU. I've launch the training and it's really slow, checking the GPU usage I've noticed that none GPU is used. There is a way to exploit the GPU?

[jonbarron]

Yep, this code is designed to run on GPUs. If your GPU is unused when you run the code, you likely have installed JAX incorrectly. I'd take a look at the JAX installation instructions for GPUs/CUDA.