Not working on cpu

[yazdaniamir38]

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]

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):