Pick a GPU device

AI
CUDA_VISIBLE_DEVICES=0 python train.py