pytorch 多卡机器中使用单卡训练

假设当前设备有两块gpu,gpu0, gpu1
仅想使用1号gpu时
1.在程序前添加
os.environ['CUDA_VISIBLE_DEVICES'] = '1' 
2. torch.load时要记得指定device 
torch.load(array, map_location=device)
3. to('cuda:0')
这里的0号gpu指的就是物理机的1号gpu


暂时只看到这些地方需要改

你可能感兴趣的:(深度学习,pytorch,计算机视觉)