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

(Jitender Singh Virk) #1

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.

0 Likes

#2

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

1 Like

(Jitender Singh Virk) #3

It worked! Thanks!

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

0 Likes