Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

pytorch when do I need to use `.to(device)` on a model or tensor?

I am new to Pytorch, but it seems pretty nice. My only question was when to use tensor.to(device) or Module.nn.to(device).

I was reading the documentation on this topic, and it indicates that this method will move the tensor or model to the specified device. But I was not clear for what operations this is necessary, and what kind of errors I will get if I don't use .to() at the right time?

For example, if I just create a tensor, I imagine that the tensor is stored in CPU accessible memory until I move the tensor to the GPU. Once the tensor is on the GPU, then the GPU will execute any mathematical operations on that tensor.

However, do I have to worry about accidentally transferring the data tensor to the GPU while not transferring the model to the GPU? Will this just give me straight errors, or will it engage in a lot of expensive data transfer behind the scenes. This example is easy enough for me to test, but I was just wondering about other cases where it might not be so obvious.

Any guidance would be helpful.

like image 302
krishnab Avatar asked Feb 03 '23 14:02

krishnab


1 Answers

It is necessary to have both the model, and the data on the same device, either CPU or GPU, for the model to process data. Data on CPU and model on GPU, or vice-versa, will result in a Runtime error.

You can set a variable device to cuda if it's available, else it will be set to cpu, and then transfer data and model to device :

import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
data = data.to(device)
like image 123
Garima Jain Avatar answered Feb 06 '23 14:02

Garima Jain