def forward(self, x):
        """Returns the output tensor.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor of dimension [B, N, S], where B is the batch size, 
            N is the number of channels, and S is the sequence length.

        Returns
        -------
        out : torch.Tensor
            Output tensor of dimension [B, spks, N, S], where spks is the number of sources 
            (speakers) and is ordered such that the first index corresponds to the target speech.
        """

        # Normalize the input
        # [B, N, L]
        x = self.norm(x)

        # Apply encoder convolution
        # [B, N, L]
        x = self.conv1d_encoder(x)

        if self.use_global_pos_enc:
            base = x  # Store the base for adding positional embedding
            x = x.transpose(1, -1)  # Change shape to [B, L, N] for positional encoding
            emb = self.pos_enc(x)  # Get positional embeddings
            emb = emb.transpose(0, -1)  # Change back to [B, N, L]
            x = base + emb  # Add positional embeddings to the base

        # Process through the computation block
        # [B, N, S]
        x = self.mdl(x)
        x = self.prelu(x)  # Apply activation

        # Expand to multiple speakers
        # [B, N*spks, S]
        x = self.conv1d_out(x)
        B, _, S = x.shape  # Unpack the batch size and sequence length

        # Reshape to [B*spks, N, S]
        # This prepares the output for gating
        # [B*spks, N, S]
        x = x.view(B * self.num_spks, -1, S)

        # Apply gated output layers
        # [B*spks, N, S]
        x = self.output(x) * self.output_gate(x)  # Element-wise multiplication for gating

        # Decode to final output
        # [B*spks, N, S]
        x = self.conv1_decoder(x)

        # Reshape to [B, spks, N, S] for output
        # [B, spks, N, S]
        _, N, L = x.shape
        x = x.view(B, self.num_spks, N, L)  # Final reshaping for output
        x = self.activation(x)  # Apply final activation

        # Transpose to [spks, B, N, S] for output
        # return the 1st spk signal as the target speech
        x = x.transpose(0, 1)
        return x[0].transpose(1, 2)  # Return only the first speaker's signal
Returns the output tensor.

        Arguments
        ---------
        x : torch.Tensor
            Input tensor of dimension [B, N, S], where B is the batch size, 
            N is the number of channels, and S is the sequence length.

        Returns
        -------
        out : torch.Tensor
            Output tensor of dimension [B, spks, N, S], where spks is the number of sources 
            (speakers) and is ordered such that the first index corresponds to the target speech.
        
clearvoice/clearvoice/models/mossformer2_se/mossformer2.py