from torchsummary import summary
summ = summary(model, torch.Size((1,28,28)))