【Pytorch】多入力モデルでテンソルのGPU移動が煩雑になってしまう

Question

Pytorchでmodelに複数の入力を与えたいときに inputs内のテンソルをGPUに移動させたいのですが、入力数に関係なく一度に行う方法が分かりません。

２入力のモデルの場合、以下のコードは煩雑ですが動作します。

for inputs, targets in train_dataloader:
    inputs[0] = inputs[0].to(device)
    inputs[1] = inputs[1].to(device)
    targets = targets.to(device)

    outputs = model(inputs)

しかし以下の方法ではスコープがfor文中のみとなり動作しません。

for inputs, targets in train_dataloader:
    for input_ in inputs:
        input_ = input_.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

そもそも以下のようなmodelの書き方が間違っているのでしょうか。

class Model(nn.Module):

    def forward(self, inputs):
        x, x1 = inputs

        x = self.model_resnet(x)
        x1 = torch.unsqueeze(x1, 1)

        features = torch.cat((x, x1), 1)
        outputs = self.fc(features)

        return outputs

入力数に関係なく一度にGPUへ移動する方法を教えていただけると幸いです。

Answer 1

Pythonの基本構文を使えばOKでした。

inputs = map(lambda input: input.to(device), inputs)