torch.device(‘cuda:0‘ if cuda else ‘cpu‘)

# 这个地方的cuda:0 实际上并不是0号GPU,他取决于CUDA_VISIBLE_DEVICES
# 然后逻辑GPU和物理GPU有一个对应关系
# 如果CUDA_VISIBLE_DEVICES为2,1,3
# 那么CUDA:0就是2号GPU, CUDA:1 就是1号GPU CUDA:3 就是3号GPU
return torch.device('cuda:0' if cuda else 'cpu')

欢迎加入大家庭:

企鹅交流群号:722058911

你可能感兴趣的:(pytorch,深度学习,神经网络)