src/transformers/models/fnet/modeling_fnet.py/FNetBasicFourierTransform/forward
class FNetBasicFourierTransform:    def forward(self, hidden_states):
        # NOTE: We do not use torch.vmap as it is not integrated into PyTorch stable versions.
        # Interested users can modify the code to use vmap from the nightly versions, getting the vmap from here:
        # https://pytorch.org/docs/master/generated/torch.vmap.html. Note that fourier transform methods will need
        # change accordingly.

        outputs = self.fourier_transform(hidden_states).real
        return (outputs,)