promach:

see the update


Regardless of the problem in the code shown above, this issue happens when one tries to convert a tensor value stored in the GPU (cuda) to numpy, if the tensor variable is called x, then, x.numpy() would raise such an error. The solution is easy by bringing back the data to the cpu and then using numpy, as follows:
x.cpu().numpy()
If, however, one is interested in finding argmax, this could be done via output.argmax() as PyTorch does support it.