Converting numpy array to tensor on GPU

import torch
from skimage import io

img = io.imread('input.png')

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

img =torch.tensor(img, device=device).float()
print(img.device)

Output

cuda:0
cpu

I am not able to convert a numpy array into a torch tensor on GPU. What is the right way to do this?

4 Likes

You should transform numpy arrays to PyTorch tensors with torch.from_numpy.
Otherwise some weird issues might occur.

img = torch.from_numpy(img).float().to(device)
19 Likes

That works. Thanks! :smile:

Thank you man this was helpful

1 Like

How to achieve this in Pytorch version 0.3.0?

You should use .cpu() or .cuda() instead of the .to(device) method. The others should also exist in 0.3.0.

7 Likes