欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

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

程序员文章站 2024-03-24 10:21:16
...
if cuda_idx >= 0:
    os.environ['CUDA_VISIBLE_DEVICES'] = str(cuda_idx) # cuda_idx在config中设置

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

相关标签: python