os.environ[‘CUDA_VISIBLE_DEVICES‘]使用指定的GPU及GPU显存

if cuda_idx >= 0:
    os.environ['CUDA_VISIBLE_DEVICES'] = str(cuda_idx) # cuda_idx在config中设置

服务器上有好几张显卡，别人可能在某几张卡上跑满了，这时需要你指定某张显卡来跑实验。使用上述功能，则代码只能对 str(cuda_idx)的显卡可见，并对所有可见的显卡依然按照数组从0开始编号。