def _init_layers(self):
        """"Init extra hidden linear layer to handle dist token if exists."""
        super(DeiTClsHead, self)._init_layers()
        if self.hidden_dim is None:
            head_dist = nn.Linear(self.in_channels, self.num_classes)
        else:
            head_dist = nn.Linear(self.hidden_dim, self.num_classes)
        self.layers.add_module('head_dist', head_dist)
"Init extra hidden linear layer to handle dist token if exists.
mmpretrain/models/heads/deit_head.py