    def forward(self, hidden_states) -> Tensor:
        """
        input hidden_states shape: [batch, height, width, channel]
        output shape: [batch, height, width, channel]
        """

        residual = hidden_states

        # norm
        hidden_states = self.group_norm(hidden_states)
        o_shape = hidden_states.shape()
        batch_dim = o_shape[0]

        hidden_states = reshape()(
            hidden_states,
            [batch_dim, -1, self.channels],
        )

        res = self.attention(hidden_states, hidden_states, hidden_states, residual) * (
            1 / self.rescale_output_factor
        )

        res = reshape()(res, o_shape)
        return res
input hidden_states shape: [batch, height, width, channel]
output shape: [batch, height, width, channel]
examples/05_stable_diffusion/src/modeling/attention.py