Since you are wrapping all submodules into an nn.Sequential
module, you are losing the flatten operation which is performed in the forward
method.
Similar to this issue you can add a custom Flatten
module and it should work:
class Flatten(nn.Module):
def __init__(self):
super(Flatten, self).__init__()
def forward(self, x):
x = x.view(x.size(0), -1)
return x
class ALEXNET(nn.Module):
def __init__(self):
super(ALEXNET, self).__init__()
self.model = models.alexnet(pretrained=False)
self.model = nn.Sequential(*(list(self.model.features.children()) + [nn.AvgPool2d(1), Flatten()] + list(self.model.classifier.children())))
def forward(self, images):
scores = self.model(images)
return scores