Tensorboard graph for sequential model with multiple inputs

tensorboard
1

The torch.utils.tensorboard.writer.SummaryWriter.add_graph method takes a input_to_model which is "a variable or a tuple of variables to be fed, and this tuple is expanded into multiples arguments.

But for torch.nn.Sequential models, the only way to have models that takes multiple arguments is to make models that takes only one argument (a tuple) and deconstruct it explicitly in the forward code (see here).

Why aren’t both behavior consistent?

As for now, if I want to make both work together, I need to give input_to_model=((x, y),) to add_graph, which is not nice.

import torch
import torch.nn as nn
from torch.utils.tensorboard.writer import SummaryWriter

class Model_X_x(nn.Module):
    def forward(self, X):
        x, y = X
        return x

class Model_x_xx(nn.Module):
    def forward(self, x):
        return x, x

x, y = torch.randn(3), torch.randn(5)
model = nn.Sequential(Model_X_x(), Model_x_xx())
pred = model((x, y))
SummaryWriter().add_graph(model, ((x, y),))