def _initialize_linear_layers(module: Module) -> None:
    """Initializes linear layers in the given module."""

    def init(mod: Module) -> None:
        if isinstance(mod, Linear):
            nn.init.xavier_uniform_(mod.weight)
            if mod.bias is not None:
                nn.init.constant_(mod.bias, 0)

    module.apply(init)
Initializes linear layers in the given module.
lightly/models/modules/masked_autoencoder.py