how to use gpu during training

google-research / multinerf

A Code Release for Mip-NeRF 360, Ref-NeRF, and RawNeRF

Apache License 2.0

3.58k stars 339 forks source link

how to use gpu during training #54

Open ZepSbosnia opened 1 year ago

ZepSbosnia commented 1 year ago

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 commented 1 year ago

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.