sksq96 / pytorch-summary

Model summary in PyTorch similar to `model.summary()` in Keras
MIT License
3.98k stars 412 forks source link

Unhelpful error message for model having single input and multiple outputs #195

Open harryknee opened 11 months ago

harryknee commented 11 months ago

I have an encoder model that takes one input and returns two values as outputs, only one is used for calculating loss. When I create a summary, the two outputs are concatenated and causes a shape issue.

E.g. a model returns both the encoded input and an intermediate representation of size [a, b, c], the torch summary would assume this to be [2, a , b, c].

It would be nice to be able to have an error message that warns user that if the model inputs and outputs size is not matching it can lead to undefined behaviour.