src/transformers/models/levit/modeling_levit.py/LevitClassificationLayer/__init__
class LevitClassificationLayer:    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.batch_norm = nn.BatchNorm1d(input_dim)
        self.linear = nn.Linear(input_dim, output_dim)