How can I change this GPU code in CPU code?

422 Views Asked by At

All.

This is my simple GPU code in PyTorch.

How can I change this in CPU code?

filename=r’./test/bees/1.jpg’
img = skimage.io.imread(filename)

x = V(centre_crop(img).unsqueeze(0), volatile=True).cuda()

model = models.dict(’resnet18’)
model = torch.nn.DataParallel(model).cuda()

model = torch.load(‘model5.pth’)

logit = model(x)
print(logit)

Thanks in advance .

1

There are 1 best solutions below

0
On

Enter this code before you call .cuda() method

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

and replacing every .cuda() with .to(device) Great answer from here