def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of IFNet.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor with shape (batch_size, n_chans, n_times).

        Returns
        -------
        torch.Tensor
            Output tensor with shape (batch_size, n_outputs).
        """
        # Pass through the spectral filtering layer
        x = self.spectral_filtering(x)
        # x is now of shape (batch_size, n_bands, n_chans, n_times)
        x = self.ensuredim(x)
        # x is now of shape (batch_size, n_bands * n_chans, n_times)

        # Pass through the feature block
        x = self.feature_block(x)

        # Flatten and pass through the final layer
        x = self.flatten(x)

        x = self.final_layer(x)

        return x
Forward pass of IFNet.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor with shape (batch_size, n_chans, n_times).

        Returns
        -------
        torch.Tensor
            Output tensor with shape (batch_size, n_outputs).
        
braindecode/models/ifnet.py