Pytorch переместит nn.module в cuda, включая подмодули ⇐ Python

[Anonymous]

Я хотел бы переместить свой модуль в cuda.

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

class SubModel(nn.Module):
def __init__(self):
super(SubModel, self).__init__()
self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=2)

def forward(self, x):
...

class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.subModel = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=2)

def forward(self, x):
...

Однако я столкнулся с ошибкой времени выполнения при выполнении self.conv1(x) в прямой функции в SubModel:
RuntimeError: Тип ввода (torch.cuda.FloatTensor ) и тип веса (torch.FloatTensor) должны быть одинаковыми
Затем я попытался переместить подмодель в cuda в конце:

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

model.subModel.to(device)

это не помогло.
Наконец, я перемещаю conv1D из подмодели в cuda:

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

class SubModel(nn.Module):
def __init__(self):
super(SubModel, self).__init__()
self.conv1 = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=2).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

Все работает отлично! Однако я чувствую, что это слишком неудобно, что мне приходится вручную перемещать все экземпляры в cuda.
Мой вопрос: существует ли удобный способ переместить модель в cuda, включая все подмодели внутри него? Спасибо!

Подробнее здесь: https://stackoverflow.com/questions/786 ... submodules