pytorch、tensorflow-多gpu显存使用问题（解决gpu 0 显存占满时 使用其他gpu 报错 out of memory问题）

问题情景：

服务器上有两个gpu，编号为0的gpu在跑任务，显存完全占满。gpu 1 空闲。在想用gpu 1 跑任务时，报错 CUDA：out of memory. 此时已经设置了 device = torch.device(“cuda:1”） ，尝试了设置device = torch.device(“cpu”），仍然会报错。

原因：

gpu 0 为主gpu，运行pytorch程序时 仍会尝试向其上写入内容，导致报错。

解决方案：

修改环境变量，让python程序忽视 gpu 0
比如 我要只使用gpu 1：

CUDA_VISIBLE_DEVICES=1 python my_script.py

如需在python程序内设置，使用如下语句：

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'

参考资料：

https://github.com/pytorch/pytorch/issues/3477