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.