def update_drop_path_rate(
    model: VisionTransformer,
    drop_path_rate: float,
    mode: str = "linear",
) -> None:
    """Updates the drop path rate in a TIMM VisionTransformer model.

    Args:
        model:
            TIMM VisionTransformer model.
        drop_path_rate:
            Maximum drop path rate.
        mode:
            Drop path rate update mode. Can be "linear" or "uniform". Linear increases
            the drop path rate from 0 to drop_path_rate over the depth of the model.
            Uniform sets the drop path rate to drop_path_rate for all blocks.
    Raises:
        ValueError: If an unknown mode is provided.
    """
    from timm.layers import DropPath

    total_depth = len(model.blocks)

    # Determine drop path rates based on the specified mode
    if mode == "linear":
        drop_probabilities = np.linspace(0, drop_path_rate, total_depth)
    elif mode == "uniform":
        drop_probabilities = [drop_path_rate for _ in range(total_depth)]
    else:
        raise ValueError(
            f"Unknown mode: '{mode}', supported modes are 'linear' and 'uniform'."
        )

    # Update the drop path rate for each block in the model
    for block, drop_prob in zip(model.blocks, drop_probabilities):
        if drop_prob > 0.0:
            block.drop_path1 = DropPath(drop_prob=drop_path_rate)
            block.drop_path2 = DropPath(drop_prob=drop_path_rate)
        else:
            block.drop_path1 = Identity()
            block.drop_path2 = Identity()
Updates the drop path rate in a TIMM VisionTransformer model.

    Args:
        model:
            TIMM VisionTransformer model.
        drop_path_rate:
            Maximum drop path rate.
        mode:
            Drop path rate update mode. Can be "linear" or "uniform". Linear increases
            the drop path rate from 0 to drop_path_rate over the depth of the model.
            Uniform sets the drop path rate to drop_path_rate for all blocks.
    Raises:
        ValueError: If an unknown mode is provided.
    
lightly/models/utils.py