hidden = torch.Tensor(1, n_hidden)
out, hidden = rnn(cat, dummy_input, hidden)
with SummaryWriter(comment="RNN") as w:
w.add_graph(rnn, (cat, dummy_input, hidden), verbose=False)
After Change
with SummaryWriter(comment="vgg19") as w:
model = torchvision.models.vgg19()
w.add_graph(model, (dummy_input, ))
with SummaryWriter(comment="densenet121") as w:
model = torchvision.models.densenet121()