TypeError: can't convert CUDA tensor to numpy.的解决办法

当试图将存储在GPU(cuda)中的张量值转换为numpy时,假设张量变量为x,然后,使用x.numpy()会引发:TypeError: can’t convert CUDA tensor to numpy,这样的错误。
其实解决方案很简单,通过将数据恢复到cpu然后使用numpy即可,如下所示:

x.cpu().numpy()

但是,如果想要找到argmax,可以通过output.argmax()来完成,因为PyTorch支持它。

猜你喜欢

转载自blog.csdn.net/qq_18649781/article/details/89214233