def forward(self, X: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.

        Parameters
        ----------
        X: Tensor
            Input tensor of shape (batch_size, n_channels, n_times).
        Returns
        -------
        Tensor
            Output tensor of shape (batch_size, n_outputs).
        """
        X = self.torch_stft(X)

        for conv in self.convs[:-1]:
            X = conv.forward(X)
        emb = self.convs[-1](X).squeeze(-1).squeeze(-1)
        return self.final_layer(emb)

        Forward pass.

        Parameters
        ----------
        X: Tensor
            Input tensor of shape (batch_size, n_channels, n_times).
        Returns
        -------
        Tensor
            Output tensor of shape (batch_size, n_outputs).
        
braindecode/models/contrawr.py