You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Describe the bug
If I try to use summary on a model that returns a list, it will only print the output shape of the first element in the list.
To Reproduce
This issue can be reproduced with a very simple model that would run like this
import torch
from torch import nn
from torchinfo import summary
class NeuralNetwork(nn.Module):
def __init__(self):
super(NeuralNetwork, self).__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28*28, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, 10),
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return [logits, logits, logits]
if __name__=="__main__":
device = "cuda" if torch.cuda.is_available() else "cpu"
model = NeuralNetwork()
print(summary(model, input_data=torch.rand(1, 28, 28, device=device), col_names = ["output_size"]))
Expected behavior
Usually models that return lists contain tensors. For instance, in a Feature Pyramid Network (FPN), one would often return a list with the different outputs for each level. Therefore, it would be nice to have a pretty way to print those shapes.
The calculate_size function should be modified accordingly for all types of input.
I remember seeing a fix for #152 to take into account modules that return numpy.ndarrays. Perhaps it would be nice to have a unified approach for different types.
Desktop (please complete the following information):
OS: Ubuntu 22.04
Pytorch 1.12.0
Torchinfo 1.7.2
The text was updated successfully, but these errors were encountered:
Describe the bug
If I try to use summary on a model that returns a list, it will only print the output shape of the first element in the list.
To Reproduce
This issue can be reproduced with a very simple model that would run like this
Expected behavior
Usually models that return lists contain tensors. For instance, in a Feature Pyramid Network (FPN), one would often return a list with the different outputs for each level. Therefore, it would be nice to have a pretty way to print those shapes.
The calculate_size function should be modified accordingly for all types of input.
Something cool would be to have:
As a reminder, version 1.7.2 would only show
I remember seeing a fix for #152 to take into account modules that return numpy.ndarrays. Perhaps it would be nice to have a unified approach for different types.
Desktop (please complete the following information):
The text was updated successfully, but these errors were encountered: