问题描述:在代码中这样写:
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