Set up single card training.
Directly in the terminal:
CUDA_VISIBLE_DEVICES=1 python train.py
Or at the beginning of the program:
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"#单卡
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2"#多卡