search

sksq96 / pytorch-summary

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

Unable to handle atomic values (int, float, etc) in output #110

Open wyc1997 opened 4 years ago

wyc1997 commented 4 years ago

When I tried to run summary() on a pytorch YOLOv3 implementation an error was thrown:

...
  File "/home/secret/anaconda3/lib/python3.7/site-packages/torchsummary/torchsummary.py", line 23, in <listcomp>
    [-1] + list(o.size())[1:] for o in output
AttributeError: 'int' object has no attribute 'size'

I looked into the both the models and the torchsummary and realize that the model implementation returns both the output tensor and the loss value as a tuple. Therefore, when the loss value is passed into the output o.size() will fail. This can be fix by checking for list and tuple before passing o to o.size()