def forward(self, x: torch.Tensor) -> torch.Tensor:
        """

        Parameters
        ----------
        X: Tensor
            Input tensor of shape (batch_size, n_channels, n_freqs, n_times).

        Returns
        -------
        Tensor
            Output tensor of shape (batch_size, n_channels, n_freqs, n_times).
        """
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.use_downsampling:
            residual = self.downsample(x)
            out += residual
        if self.pooling:
            out = self.maxpool(out)
        out = self.dropout(out)
        return out


        Parameters
        ----------
        X: Tensor
            Input tensor of shape (batch_size, n_channels, n_freqs, n_times).

        Returns
        -------
        Tensor
            Output tensor of shape (batch_size, n_channels, n_freqs, n_times).
        
braindecode/models/contrawr.py