def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        r"""Forward method of the `FourierFeatures` class."""
        original_dtype = inputs.dtype
        inputs = inputs.to(torch.float32)
        num_channels = inputs.shape[1]
        num_freqs = (self.stop - self.start) // self.step

        freqs = torch.arange(self.start, self.stop, self.step, dtype=inputs.dtype, device=inputs.device)
        w = torch.pow(2.0, freqs) * (2 * torch.pi)  # [num_freqs]
        w = w.repeat(num_channels)[None, :, None, None, None]  # [1, num_channels * num_freqs, 1, 1, 1]

        # Interleaved repeat of input channels to match w
        h = inputs.repeat_interleave(
            num_freqs, dim=1, output_size=inputs.shape[1] * num_freqs
        )  # [B, C * num_freqs, T, H, W]
        # Scale channels by frequency.
        h = w * h

        return torch.cat([inputs, torch.sin(h), torch.cos(h)], dim=1).to(original_dtype)
Forward method of the `FourierFeatures` class.
src/diffusers/models/autoencoders/autoencoder_kl_mochi.py