search

sksq96 / pytorch-summary

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

bug: in case an axis=-1 #160

Open hanochk opened 3 years ago

hanochk commented 3 years ago

https://github.com/sksq96/pytorch-summary/blob/011b2bd0ec7153d5842c1b37d1944fc6a7bf5feb/torchsummary/torchsummary.py#L94

Rplace with np.abs(np.prod(summary[layer]["output_shape"])) if in case the summary[layer]["output_shape"] = [-1, 32, 128, 128]