If I call model.cuda() in pytorch where model is a subclass of nn.Module, and say if I have four GPUs, how it will utilize the four GPUs and how do I know which GPUs that are using?

  • $\begingroup$ Please consider marking the answer as correct, or alternatively commenting why you think it is not correct. $\endgroup$
    – noe
    Feb 16 '21 at 22:49

model.cuda() by default will send your model to the "current device", which can be set with torch.cuda.set_device(device).

An alternative way to send the model to a specific device is model.to(torch.device('cuda:0')).

This, of course, is subject to the device visibility specified in the environment variable CUDA_VISIBLE_DEVICES.

You can check GPU usage with nvidia-smi. Also, nvtop is very nice for this.

The standard way in PyTorch to train a model in multiple GPUs is to use nn.DataParallel which copies the model to the GPUs and during training splits the batch among them and combines the individual outputs.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.