def __init__(self, h, d_model, after_reduced_cnn_size, dropout=0.1):
        """Take in model size and number of heads."""
        super().__init__()
        assert d_model % h == 0
        self.d_per_head = d_model // h
        self.h = h

        base_conv = CausalConv1d(
            in_channels=after_reduced_cnn_size,
            out_channels=after_reduced_cnn_size,
            kernel_size=7,
            stride=1,
        )
        self.convs = nn.ModuleList([deepcopy(base_conv) for _ in range(3)])

        self.linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(p=dropout)
Take in model size and number of heads.
braindecode/models/sleep_stager_eldele_2021.py