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

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

        Returns
        -------
        torch.Tensor
            Output tensor of shape (batch_size, n_outputs).
        """
        # x shape: (batch_size, n_chans, n_times)
        x = self.block_1(x)
        # (batch_size, n_filter, n_chans, n_times)
        x = self.block_2(x)
        # (batch_size, n_filter*4, n_chans, n_times)
        x = self.block_3(x)
        # (batch_size, 1, n_filter*8, n_times//4)
        x = self.block_4(x)
        # (batch_size, 1, n_filter*8, n_times//4)
        x = self.block_5(x)
        # (batch_size, n_filter*(n_times//32))
        x = self.final_layer(x)

        return x

        Forward pass of the EEGNeX model.

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

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