Hello everybody,
I am trying to deploy different pytorch-based training scripts on different GPUs. However, the information I could find is about training a model on multiple GPUs.
Could some tell me how to do this?
I tried the ‘spawn’ trick, ‘.cuda(0)/.cuda(1)’ trick… But they were not working.
Sorry if this is a bad question
Thank you!
just need to pass CUDA_VISIBLE_DEVICES=0 python script.py
replace 0 by the gpu you want
1 Like
You could pass the device you want to train on as an argument to the script.
For example, ‘cuda:0’ corresponds to the 1st GPU in your system, ‘cuda:1’ corresponds to the 2nd GPU and so on.
Then assuming you store the passed argument in a variable named device
, all you have to do is to call .to(device)
on your tensors etc.
1 Like