How to use with torch.cuda.device() conditionally

user_4391249 0
Viewed 1600

I have some code encapsulated in a:

with torch.cuda.device(self.device):
    # do a bunch of stuff

And in my __init__ I have:

self.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

But I'm a little confused about how to deal with a situation where the device is cpu. Because torch.cuda.device is already explicitly for cuda. Should I just write a decorator for the function? Seems a bit overkill

pytorch
1 Answers

According to the documentation for torch.cuda.device

device (torch.device or int) – device index to select. It’s a no-op if this argument is a negative integer or None.

Based on that we could use something like

with torch.cuda.device(self.device if self.device.type == 'cuda' else None):
    # do a bunch of stuff

which would simply be a no-op when self.device isn't a CUDA device.

user_27900470
Related