Hi,
I am trying to create a custom nn module that wraps a pre-trained densenet121, so that I can access activations at the end of the convolutional part of the network.
As a starting point, I have just done the following:
import torch
from torchvision import models
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.model = models.densenet121(pretrained=True)

        self.features = self.model.features
        self.classifier = self.model.classifier

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

example = torch.ones((1,3,224,224))

model = MyModel()
output = model.forward(example)

However,  I get a mismatch error for line x = self.classifier(x) of the forward method:
RuntimeError: size mismatch, m1: [7168 x 7], m2: [1024 x 1000] at ../aten/src/TH/generic/THTensorMath.cpp:961

Also, as far as I can tell, densenet121 expects (batch_size, 1024) inputs to its classifier part, however when I print the size of x after x = self.features(x) in the forward method, it is (1, 1024, 7, 7)?
Thank you for the help, I am still very new to pytorch.