sksq96 / pytorch-summary

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

support for list as output item, not only tensors #167

Open leora-betesh opened 3 years ago

leora-betesh commented 3 years ago

Add support for a net outputting a list. Say you are returning intermediary feature-maps from a U-net. These feature-maps will have a different size and cannot be stacked into one tensor. Instead they can be placed in a list.