how to fine-tuning the cspnet model?

[GinkgoX]

Dear rwightman:
thanks for your great work for the pytorch integration of image classification models, when I try to use the cspnet to fine-tune my individal dataset, but the error still exists, the error info as follows:
(where I have changed the head as: self.model.head = nn.Linear(2048, n_class))

~/pytorch-image-models-master/timm/models/cspnet.py in forward(self, x)
401 def forward(self, x):
402 x = self.forward_features(x)
--> 403 x = self.head(x)
404 return x
405

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725 result = self._slow_forward(*input, **kwargs)
726 else:
--> 727 result = self.forward(*input, **kwargs)
728 for hook in itertools.chain(
729 _global_forward_hooks.values(),

~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/linear.py in forward(self, input)
91
92 def forward(self, input: Tensor) -> Tensor:
---> 93 return F.linear(input, self.weight, self.bias)
94
95 def extra_repr(self) -> str:

~/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py in linear(input, weight, bias)
1690 ret = torch.addmm(bias, input, weight.t())
1691 else:
-> 1692 output = input.matmul(weight.t())
1693 if bias is not None:
1694 output += bias

RuntimeError: mat1 dim 1 must match mat2 dim 0

could you help me about this error, thanks a lot !!!

[rwightman]

@GinkgoX you might want to look at the code of the model and see what the head looks like, it's not just the linear layer.

If you create the model with the factory it will replace the layer correctly for you timm.create_model('cspresnet50', num_classes=n_class, pretrained=True)

[GinkgoX]

thanks for your response, I have tried this, but in inference stage, I still get error, which indicates that the expect dim is n_class, but get 1000, so I print the model, and change the head as follows:
n_features = self.model.head.fc.in_features
self.model.head.fc = nn.Linear(n_features, n_class)
it runs ok for both training and test stage. but still appreciate your help ~ \ ( ^ v ^ ) /