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