os.environ[“CUDA_VISIBLE_DEVICES“] = args.device不起作用

问题描述：在代码中这样写：

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
    main()

不起作用。原因是os.environ["CUDA_VISIBLE_DEVICES"] = args.device 必须要处于所有使用cuda的代码之前才能生效。然而我的全局变量中有一个 stack_kernel = torch.tensor(stack_kernel).cuda()，所以在main函数执行之前就是用了cuda.

解决方案：在命令行中使用 os.environ["CUDA_VISIBLE_DEVICES"] = device_id

CUDA_VISIBLE_DEVICES=device_id python your_program.py
版权声明：本文为xjtu_yyl原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接和本声明。