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

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

I found out that I can just use

output.argmax()

numpy neural-network pytorch pruning

You can use torch.max function like follows:

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

CodeHunter