sksq96 / pytorch-summary

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

Bug if number of parameters is zero. #183

Open levifussell opened 2 years ago

levifussell commented 2 years ago

Hi,

If the number of parameters of the model is zero, then the call for (line 102):

total_params_size = abs(total_params.numpy() * 4. / (1024 ** 2.))

will fail because the type is an integer, not a torch tensor. This may seem like a weird case (when are there no parameters), but I came across this after writing 'wrapper' Modules for modules with parameters. These wrappers had no parameters themselves.

The fix is just to check the parameter count is non-zero, otherwise remove the 'numpy()' call.