def forward(self, x):
        """
        Defines the forward pass of the FSMN_Wrap module.

        Args:
            x (torch.Tensor): Input tensor of shape [batch_size, channels, height, time, 2].

        Returns:
            torch.Tensor: Output tensor reshaped to [batch_size, channels, height, time].
        """
        # Shape of input x: [b, c, h, T, 2]
        b, c, T, h = x.size()

        # Permute x to reshape it for FSMN processing: [b, T, h, c]
        x = x.permute(0, 2, 3, 1)  # Change dimensions to [b, T, h, c]
        x = torch.reshape(x, (b * T, h, c))  # Reshape to [b*T, h, c]

        # Pass through the FSMN
        output = self.fsmn(x)  # output: [b*T, h, c]

        # Reshape output back to original dimensions
        output = torch.reshape(output, (b, T, h, c))  # output: [b, T, h, c]

        return output.permute(0, 3, 1, 2)  # Final output shape: [b, c, h, T]

        Defines the forward pass of the FSMN_Wrap module.

        Args:
            x (torch.Tensor): Input tensor of shape [batch_size, channels, height, time, 2].

        Returns:
            torch.Tensor: Output tensor reshaped to [batch_size, channels, height, time].
        
clearvoice/clearvoice/models/mossformer_gan_se/generator.py