pytorch 指定GPU训练
指定GPU代码
上述使用中，代码3最为常用，keras也可以通过这种方式来调用GPU训练神经网络。
# 代码1：
torch.cuda.set_device(1)
# 代码2：
device = torch.device("cuda:1")
# 代码3：（官方推荐使用）
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
（如果你想同时调用两块GPU的话）
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
使用过程中遇到的问题
问题：在使用上述代码过程中，程序仍然使用默认GPU训练。
如图：
之后，程序报错：
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
那么，你可能需要将指定GPU代码放在程序段最开始的部位。例如：