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

[Anonymous]

Я новичок в использовании Pytorch. Я хотел бы переместить свой модуль в cuda.

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

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

def forward(self, x):
...
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.submodel = SubModel()
self.conv = nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3)

def forward(self, x):
...

model = MyModel()
model.to(device)

Однако я столкнулся с ошибкой во время выполнения: RuntimeError: тип ввода (torch.cuda.FloatTensor) и тип веса (torch.FloatTensor) должны быть одинаковыми< /p>
Затем я попытался переместить подмодель в 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=3),to(torch.device('cuda'))

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

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