Как удобно получить тип устройства модуля pytorch? ⇐ Python

[Anonymous]

Мне нужно наложить несколько собственных слоев на разные типы моделей Pytorch и на разных устройствах.
Например. A — это модель cuda, а B — модель процессора (но я не знаю ее, пока не узнаю тип устройства). Тогда новыми моделями будут C и D соответственно, где

Код: Выделить всё

class NewModule(torch.nn.Module):
def __init__(self, base):
super(NewModule, self).__init__()
self.base = base
self.extra = my_layer() # e.g. torch.nn.Linear()

def forward(self,x):
y = self.base(x)
z = self.extra(y)
return z

...

C = NewModule(A) # cuda
D = NewModule(B) # cpu

Однако я должен переместить базовый и дополнительный на то же самое устройство, т. е. базовый и дополнительный C — это модели cuda, а D — модели cpu. Итак, я попробовал это __init__:

Код: Выделить всё

def __init__(self, base):
super(NewModule, self).__init__()
self.base = base
self.extra = my_layer().to(base.device)

К сожалению, в torch.nn.Module нет атрибута device (поднимите AttributeError).
Что мне следует сделать, чтобы получить базовый тип устройства? Или какой-либо другой способ заставить базу и дополнительную автоматически находиться на одном устройстве, даже если структура базы неопределенна?