Understand Difference torch.device(“cuda”) and torch.device(“cuda:0”) – PyTorch Tutorial

In python, we often need to move a tensor or model to cpu or gpu device. For example:

device = None
torch.manual_seed(h.seed)
    
if torch.cuda.is_available():
    torch.cuda.manual_seed(h.seed)
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

generator = Generator(h).to(device)
mpd = MultiPeriodDiscriminator().to(device)
msd = MultiScaleDiscriminator().to(device)

However, we may see this kind of code:

device = torch.device('cuda')
device = torch.device('cuda:0')

In this tutorial, we will introduce their difference.

As to

device = torch.device('cuda')

It will use gpu by torch.cuda.current_device(). The default value of it is 0.

As to

device = torch.device('cuda:0')

Here 0 is the index of gpu you plan to use. We can use torch.cuda.device_count() to get all gpu index.

For example:

device = torch.device('cuda:1')

It will use the second gpu in your computer.