def get_layer_id_for_vit(name, num_layers):
    """
    Assign a parameter with its layer id
    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
    """
    if name in ['encoder.cls_token', 'encoder.encoder_pos_embedding']:
        return 0
    if name.startswith('encoder.stem'):
        return 0
    if name.startswith('encoder.encoder.blocks'):
        return int(name.split('.')[3]) + 1
    return num_layers