Limit GPU cores for training

🤗Transformers
#1

My company has a 6-gpu server. However, I only want to train the model on a specific number of gpus (maybe 1 or 2).
How can I achieve that?

#2

What are you using for training ? Trainer or the pytorch-lightning ?

#3

You can set this env variable

CUDA_VISIBLE_DEVICES=0,1

specify the devices that you want to use and the Trainer will only use those cuda devices.

For pytorch-lightninig here’s official doc

1 Like
#4

Thank you. This is exactly what I need