Select GPU of choice - Githubissues

lopuhin / transformer-lm

Transformer language model (GPT-2) with sentencepiece tokenizer

164 stars 47 forks source link

Select GPU of choice #25

Closed Meghana-Meghana closed 4 years ago

Meghana-Meghana commented 4 years ago

Assuming there are 4 GPUs, how do I train on the GPU of choice instead of all the GPUs? I tried fiddling around with n_devices and device_id variables in main.py, but in vain.

lopuhin commented 4 years ago

@Meghana-Meghana you can set environment variable CUDA_VISIBLE_DEVICES, e.g. CUDA_VISIBLE_DEVICES=0 gpt-2 ... or to use two GPUs CUDA_VISIBLE_DEVICES=0,1 gpt-2 ...