def add_stochastic_depth_to_blocks(vit: Module, prob: float = 0.0, mode="row") -> None:
    """Adds stochastic depth dropout to all transformer blocks in a Vision Transformer Model

    Args:
        vit:
            Vision Transformer Model to which stochastic depth dropout will be added.
        prob:
            Probability of dropping a layer.
        mode:
            Mode for stochastic depth. Default is "row".

    Raises:
        Runtime Error: If torchvision version is less than 0.12.
    """
    if dependency.torchvision_vit_available():
        # Requires torchvision >=0.12
        from torchvision.models.vision_transformer import EncoderBlock
    else:
        raise RuntimeError("add_stochastic_depth_to_blocks requires torchvision>=0.12.")

    if prob <= 0:
        return

    for mod in vit.modules():
        if isinstance(mod, EncoderBlock):
            mod.dropout = Sequential(mod.dropout, StochasticDepth(p=prob, mode=mode))
            mod.mlp = Sequential(mod.mlp, StochasticDepth(p=prob, mode=mode))
Adds stochastic depth dropout to all transformer blocks in a Vision Transformer Model

    Args:
        vit:
            Vision Transformer Model to which stochastic depth dropout will be added.
        prob:
            Probability of dropping a layer.
        mode:
            Mode for stochastic depth. Default is "row".

    Raises:
        Runtime Error: If torchvision version is less than 0.12.
    
lightly/models/utils.py