def _init_weights(layer):
        """
        Initialize the weights of the model for each layer layer.

        If the layer is a linear layer, the weight will be initialized
        with a truncated normal distribution with std=0.02.

        If m.bias is not None, the bias will be initialized with a constant
        value of 0.

        If the layer is a layer normalization layer, the bias will be
        initialized with a constant value of 0, and the weight will be
        initialized with a constant value of 1.

        Parameters
        ----------
        m : torch.nn.Module
            The layer of the pytorch model
        """

        if isinstance(layer, nn.Linear):
            trunc_normal_(layer.weight, std=0.02)
            if layer.bias is not None:
                nn.init.constant_(layer.bias, 0)
        elif isinstance(layer, nn.LayerNorm):
            nn.init.constant_(layer.bias, 0)
            nn.init.constant_(layer.weight, 1.0)

        Initialize the weights of the model for each layer layer.

        If the layer is a linear layer, the weight will be initialized
        with a truncated normal distribution with std=0.02.

        If m.bias is not None, the bias will be initialized with a constant
        value of 0.

        If the layer is a layer normalization layer, the bias will be
        initialized with a constant value of 0, and the weight will be
        initialized with a constant value of 1.

        Parameters
        ----------
        m : torch.nn.Module
            The layer of the pytorch model
        
braindecode/models/labram.py