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() to index = output.cpu().data.numpy().argmax() This means data is first moved to cpu and then converted to numpy array.