Printing model summary

vision
#1

I am trying to print the no. of parameters in my model.

I used

!pip install pytorch-model-summary
from pytorch_model_summary import summary

batch_size_train = 2
batch_size_validate = 1

num_epochs = 10
learning_rate = 0.000009
img_width = 720
img_height = 720

model = Att_Unet(in_channels =1, out_channels =1)
model.cuda()

print(summary(model, torch.zeros(2, 1, 720, 720), show_input=True))

criterion = Diceloss().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

I get an error in google collab

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

but when I use

torch.zeros(2, 1, 720, 720).cuda()

I get error as

RuntimeError: CUDA out of memory. Tried to allocate 252.00 MiB (GPU 0; 7.43 GiB total capacity; 6.76 GiB already allocated; 74.94 MiB free; 6.84 GiB reserved in total by PyTorch)

but my model still runs without any error of CUDA out of memory, this only happens when I am trying to print model summary