sksq96 / pytorch-summary

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

Not working on cpu #179

Open yazdaniamir38 opened 2 years ago

yazdaniamir38 commented 2 years ago

Hi, I tried to use get my model summary on cpu by:

summary(model, input_size=(1,200, 200),device=torch.device('cpu'))

But I get the following error: AttributeError: 'torch.device' object has no attribute 'lower'.

My troch version is 1.10.0.

gagank1 commented 2 years ago

I'm seeing the same issue. It looks like there's a discrepancy between the version of torchsummary on pip and the version on this repo. The one on pip expects a string.

This is the function signature for summary from pip: def summary(model, input_size, batch_size=-1, device="cuda"):

and this is the one on this repo: def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None):