Check classification head

[robertdvdk]

The classifier seems like it's underperforming relative to AIM. We should check carefully whether the implementation is actually the same.