在代码里

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

或者在terminal

export CUDA_VISIBLE_DEVICES=0