Как использовать с torch.cuda.device () условно

Question

У меня есть код, заключенный в:

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

И в моем __init__ у меня есть:

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

Но я немного запутался в том, что делать с ситуация, когда на устройстве стоит процессор. Потому что torch.cuda.device уже явно для cuda. Должен ли я просто написать декоратор для функции? Кажется, немного перебор

Answer 1

Согласно документации для torch.cuda.device

device (torch.device или int) - индекс устройства для выбора. Это не работает, если этот аргумент является отрицательным целым числом или None.

Исходя из этого, мы могли бы использовать что-то вроде

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

, которое просто не работает, если self.device не является устройством CUDA.