Setting specific device for Trainer

A solution that works for me:

import os
os.environ[“CUDA_DEVICE_ORDER”]=“PCI_BUS_ID”
os.environ[“CUDA_VISIBLE_DEVICES”]=“0”

But import this FIRST, before anything from the pytorch or transformer libraries.

7 Likes