The feature dimension looks reasonable as the mlp_head would not be applied as seen here. I haven’t checked the expected shape as the einops reductions don’t show the actual values, but I would expect to see this or a similar shape.
Could you explain what the issue is or why this shape would not be expected?
Your model seems to be overfitting on the training set and I don’t think that your feature extraction is necessarily wrong.
Overfitting can have different reasons, e.g. the model capacity might be too large for the given data and your model is thus able to easily learn all training samples.
It depends which features you want to use. Initially, you’ve replaced the entire mlp_head with an nn.Identity layer, now you are using the Reduce layer. Both sound reasonable as they are applied before the final linear layer which would act as the classifier.