What torch.cuda.set_device() does?

1

Does torch.cuda.set_device(args.gpu) set a GPU for execution or it sets the number of GPUs should be used for execution?

If it sets the GPU for execution, how can I set multiple GPUs to run my experiment? For example, I want to tell to pytorch that you should use two GPUs (if available) to run my experiment. How can I achieve that?

2

torch.cuda.set_device sets the default GPU and in order to use multi-GPU, you may use nn.DataParallel

3

Can you give an example in order to use multi-GPU?

4

you may see the following example OpenNMT

1 Like