class MochiModulatedRMSNorm(nn.Module):
    def __init__(self, eps: float):
        super().__init__()

        self.eps = eps
        self.norm = RMSNorm(0, eps, False)

    def forward(self, hidden_states, scale=None):
        hidden_states_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)

        hidden_states = self.norm(hidden_states)

        if scale is not None:
            hidden_states = hidden_states * scale

        hidden_states = hidden_states.to(hidden_states_dtype)

        return hidden_states
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/transformers/transformer_mochi.py