PyTorch CUDA to Numpy?

dataiter = iter(testloader)
images, labels =

move model inputs to cuda, if GPU available

if train_on_gpu:
images = images.cuda()

get sample outputs

output = model(images)

convert output probabilities to predicted class

_, preds_tensor = torch.max(output, 1)
preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())

plot images in the batch along with the predicted and true labels

fig = plt.figure(figsize=(25, 4))
for idx in np.arange(20):
ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
ax.set_title("{} ({})".format(classes[preds[idx]], classes[labels[idx]]),
color=( ‘green’ if preds[idx] == labels[idx].item() else ‘red’))

This shows the following error while running the code.
TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

It’s weird. You have a condition that checks if the train_on_gpu is True or False. So, this error should not happen. Can you check the device attribute of the tensor preds_tensor?