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
just need to pass CUDA_VISIBLE_DEVICES=0 python script.py
replace 0 by the gpu you want
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.