def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Returns the output tensor.

        Args:
            x (torch.Tensor): Input tensor of dimension [B, N, S], where:
                B = Batch size,
                N = Number of channels (filters),
                S = Sequence length.

        Returns:
            out (torch.Tensor): Output tensor of dimension [spks, B, N, S], 
                                where:
                spks = Number of speakers,
                B = Batch size,
                N = Number of filters,
                S = Number of time frames.
        """
        # [B, N, L] - Normalize the input tensor
        x = self.norm(x)

        # [B, N, L] - Apply 1D convolution to encode features
        x = self.conv1d_encoder(x)
        
        # If using global positional encoding, add the positional embeddings
        if self.use_global_pos_enc:
            base = x  # Store the original encoded features
            x = x.transpose(1, -1)  # Change shape to [B, L, N]
            emb = self.pos_enc(x)  # Get positional embeddings
            emb = emb.transpose(0, -1)  # Change shape back to [B, N, L]
            x = base + emb  # Add positional embeddings to encoded features

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

        # [B, N*spks, S] - Project features to multiple speaker outputs
        x = self.conv1d_out(x)
        B, _, S = x.shape  # Get the shape after convolution

        # [B*spks, N, S] - Reshape for speaker outputs
        x = x.view(B * self.num_spks, -1, S)

        # [B*spks, N, S] - Apply gated output layer
        x = self.output(x) * self.output_gate(x)

        # [B*spks, N, S] - Decode back to original channel size
        x = self.conv1_decoder(x)

        # [B, spks, N, S] - Reshape output tensor to include speaker dimension
        _, N, L = x.shape
        x = x.view(B, self.num_spks, N, L)
        x = self.activation(x)  # Apply ReLU activation

        # [spks, B, N, S] - Transpose to match output format
        x = x.transpose(0, 1)

        return x  # Return the output tensor
Returns the output tensor.

        Args:
            x (torch.Tensor): Input tensor of dimension [B, N, S], where:
                B = Batch size,
                N = Number of channels (filters),
                S = Sequence length.

        Returns:
            out (torch.Tensor): Output tensor of dimension [spks, B, N, S], 
                                where:
                spks = Number of speakers,
                B = Batch size,
                N = Number of filters,
                S = Number of time frames.
        
clearvoice/clearvoice/models/mossformer2_ss/mossformer2.py