Yes you can do that. You likely forgot to merge some dimensions before giving them to the Linear layer (try print(x.size()) before self.top(x) - it will be 4D). Just do self.top(x.view(-1, 512)).
Additionally, in case of ResNet, you could just replace the final layer with a new one and it should work I think:
model_imagenet.fc = nn.Linear(512, num_classes)