How to set gpu device for hugging trainer?

🤗Transformers
1

I am trying to set gpu device for HF trainer. By default it uses device:0. I want to use a custom device. I tried using torch.cuda.set_device(). It does not work.