The error is because the model is on GPU while your input image is on CPU. We have to make sure that they are both on GPU or CPU.
@@ -107,7 +107,7 @@ def predict(net, item):
with torch.no_grad():
if torch.cuda.is_available():
- inputs_test = torch.cuda.FloatTensor(sample["image"].unsqueeze(0).float())
+ inputs_test = torch.cuda.FloatTensor(sample["image"].unsqueeze(0).cuda().float())
else:
inputs_test = torch.FloatTensor(sample["image"].unsqueeze(0).float())