如果你想在运行脚本时直接指定GPU,可以使用环境变量CUDA_VISIBLE_DEVICES来控制哪些GPU对CUDA库可见。例如,如果你想要使用第二块GPU(索引为1),你可以在运行脚本之前设置该环境变量:
export CUDA_VISIBLE_DEVICES=1
python xxx.py
这样,CUDA相关的库和框架就只会看到并使用索引为1的GPU,而忽略其他GPU。请注意,这种方法会影响在同一终端会话中运行的所有后续CUDA相关的程序,直到你更改或清除CUDA_VISIBLE_DEVICES环境变量。
在某些情况下,你可能需要重新启动终端会话以使环境变量更改生效。