I have been trying to build a model using models.densenet161 but I’m having trouble training it.
# %%capture
if not debug:
    model = models.densenet161(pretrained=True)
    # Freeze all layers
    for param in model.parameters():
        param.requires_grad = False

    model.classifier = nn.Sequential(nn.Linear(2208, 256),
                                     nn.ReLU(),
                                     nn.Dropout(0.2),
                                     nn.Linear(256, len(trainloader.dataset.classes)),
                                     nn.LogSoftmax(dim=1))

    criterion = nn.NLLLoss()
    optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)
    model.to(device);

When I try to print the summary for this model using
torchsummary.summary(model, (3, 224, 224))
I get the error
AttributeError: 'list' object has no attribute 'size'
Similarly, the same error is thrown during training at this line
logps = model.forward(inputs)
This was working completely fine with pre-trained resnet50 and I saw in other topics that to modify the network it is necessary to use classifier instead of fc. I would appreciate if someone could tell me why this error is happening.
FYI this is the internal code that breaks:
site-packages/torchsummary/torchsummary.py in hook(module, input, output)
     17             m_key = "%s-%i" % (class_name, module_idx + 1)
     18             summary[m_key] = OrderedDict()
---> 19             summary[m_key]["input_shape"] = list(input[0].size())
     20             summary[m_key]["input_shape"][0] = batch_size
     21             if isinstance(output, (list, tuple)):```

Edit 1: changed 1024 to 2208, in the first linear layer.