I a add x = x.view(x.numel()) to flatten the intermediate feature representation before the classifier part, then the dimensions match but only for dataloader batch size = 1.
    def forward(self, x):
        featuremaps = []
        for id, module in self.submodule_features._modules.items():
            x = module(x)
            if int(id) in self.extract_features_layers:
                featuremaps.append(x)
        x = x.view(x.numel())
        for id, module in self.submodule_classifier._modules.items():
            x = module(x)
            if int(id) in self.extract_classifier_layers:
                featuremaps.append(x)
        return featuremaps