src/transformers/models/levit/modeling_levit.py/MLPLayerWithBN/__init__
class MLPLayerWithBN:    def __init__(self, input_dim, output_dim, bn_weight_init=1):
        super().__init__()
        self.linear = nn.Linear(in_features=input_dim, out_features=output_dim, bias=False)
        self.batch_norm = nn.BatchNorm1d(output_dim)