Reputation: 1029

TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first

I am using a modified predict.py for testing a pruned SqueezeNet Model

[phung@archlinux SqueezeNet-Pruning]$ python predict.py --image 3_100.jpg --model model_prunned --num_class 2
prediction in progress
Traceback (most recent call last):
File “predict.py”, line 66, in
prediction = predict_image(imagepath)
File “predict.py”, line 52, in predict_image
index = output.data.numpy().argmax()
TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
[phung@archlinux SqueezeNet-Pruning]$

I understand that numpy does not support GPU(CUDA) yet.

How shall I modify the code to get away from this error without invoking tensor copy data operation, tensor.cpu()?

Upvotes: 63

Answers (3)

Umang Gupta

Reputation: 16450

Change

index = output.data.numpy().argmax()

index = output.cpu().data.numpy().argmax()

This means data is first moved to cpu and then converted to numpy array.

Upvotes: 88

Kaushik Roy

Reputation: 1685

You can use torch.max function like follows:

value, index = torch.max(output,1)

Upvotes: 2

kevin998x

Reputation: 1029

I found out that I can just use

output.argmax()

Upvotes: 5

TypeError: can’t convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first

Answers (3)

Related Questions