sksq96 / pytorch-summary

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

Added OrderedDict as an output datatype #184

Open SoumyadeepB opened 2 years ago

SoumyadeepB commented 2 years ago

Some pre-trained models like the deeplabv3_resnet50 returns an OrderedDict containing the semantic mask and the auxillary loss as two separate Tensors under the keys 'out' and 'aux'.

Usage of torchsummary with such models was not compatible since OrderedDict had not been considered as a potential output type.

Error thrown: 'collections.OrderedDict' object has no attribute 'size'

This issue has now been fixed.