def forward(self, x):
        """
        Forward pass for the OffsetScale layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            List[torch.Tensor]: A list of tensors with applied offsets and scales for each head.
        """
        out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta  # Apply scaling and offsets
        return out.unbind(dim=-2)  # Unbind heads into a list

        Forward pass for the OffsetScale layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            List[torch.Tensor]: A list of tensors with applied offsets and scales for each head.
        
train/speech_super_resolution/models/mossformer2_sr/mossformer2_block.py