def METHOD_NAME(self, x):
    """Forward pass.
    Parameters
    ----------
    x: torch.Tensor
        Batch of EEG windows of shape (batch_size, n_channels, n_times).
    """
    if x.ndim == 3:
        x = x.unsqueeze(2)
    feats = self.feature_extractor(x).flatten(start_dim=1)
    if self.return_feats:
        return feats
    else:
        return self.fc(feats)