How to specify gpu like pytorch.device('cuda:2') in fastai?

I’ve 3 gpus system in my university and i’m allowed to use 3rd gpu. In pytorch, I can specify it with torch.device(‘cuda:2’). How to do the same with fastai. Fastai automatically put code on cuda:0.

If you use torch.cuda.set_device(2) everything should be on that GPU.

2 Likes

It worked! Thanks!

torch.cuda.set_device('cuda:2')

1 Like

after we
torch.cuda.set_device('cuda:1')
when we construct a learner with dls as arguments,
we also have to

dls.device = default_device()

:grinning:

Including this line at the beginning of my code (according to the fastai documentation) worked for me:

import os; os.environ['CUDA_VISIBLE_DEVICES']='2'