How to use torch.utils.tensorboard's SummaryWriter add_graph with dictionary outputs?

Cross Posting from Stack Overflow for visibility:

I have a PyTorch model that inherits from nn.Module and that has a forward method which returns a dictionary containing multiple tensors (stored as values). When I try to log this model to TensorBoard using torch.utils.tensorboard's SummaryWriter.add_graph() method I receive the following error:

RuntimeError("Tracer cannot infer type of <output dictionary>

How can I use add_graph() when the model’s forward() method returns a dictionary?

1 Like