Setting specific device for Trainer

Beginners
24

This works for me:

import os
os.environ["CUDA_VISIBLE_DEVICES"]="1"
import torch

You have to set the env variable before importing torch. Now torch only sees device 1, not 0.

If you run this:

print(torch.cuda.current_device())
print(torch.cuda.is_available())

Torch would still think it is using device 0 but from the command nvidia-smi, it is using device 1.

5 Likes