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

Change

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

to

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

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