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.

1 Like

It worked! Thanks!


1 Like

after we
when we construct a learner with dls as arguments,
we also have to

dls.device = default_device()


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'