prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  fourier_mix attention does not support mixed precision training
  Fourier mix attention does not support half precision.
  
  ```Traceback (most recent call last):
   ...
    File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
      return forward_call(*input, **kwargs)
    File "/content/muzik/src/modules/xformer/custom_block_factory.py", line 196, in forward
      x = self.mha(q, k, v, att_mask=att_mask)
    File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
      return forward_call(*input, **kwargs)
    File "/content/muzik/src/modules/xformer/custom_multi_head_dispatch.py", line 254, in forward
      y = self.attention(q=q, k=k, v=v, att_mask=att_mask)
    File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
      return forward_call(*input, **kwargs)
    File "/usr/local/lib/python3.7/dist-packages/xformers/components/attention/fourier_mix.py", line 25, in forward
      att = torch.fft.fft2(q).real
  RuntimeError: Unsupported dtype Half```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `xformers/components/attention/fourier_mix.py`
  Content:
  ```
  1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
  2 #
  3 # This source code is licensed under the BSD license found in the
  4 # LICENSE file in the root directory of this source tree.
  5 
  6 import torch
  7 
  8 from xformers.components.attention import Attention, AttentionConfig, register_attention
  9 
  10 
  11 @register_attention("fourier_mix", AttentionConfig)
  12 class FourierMix(Attention):
  13     def __init__(self, dropout: float, *_, **__):
  14         """
  15         FFT-based pseudo-attention mechanism, from
  16         "
  17         "FNet: Mixing Tokens with Fourier Transforms"
  18         Lee-Thorp et al., 2021, https://arxiv.org/pdf/2105.03824.pdf
  19         """
  20         super().__init__()
  21         self.attn_drop = torch.nn.Dropout(dropout, inplace=False)
  22         self.requires_input_projection = False
  23 
  24     def forward(self, q: torch.Tensor, *_, **__):
  25         att = torch.fft.fft2(q).real
  26 
  27         att = self.attn_drop(att)
  28 
  29         return att
  30
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/xformers/components/attention/fourier_mix.py b/xformers/components/attention/fourier_mix.py
  --- a/xformers/components/attention/fourier_mix.py
  +++ b/xformers/components/attention/fourier_mix.py
  @@ -4,6 +4,7 @@
   # LICENSE file in the root directory of this source tree.
   
   import torch
  +from torch.cuda.amp import autocast
   
   from xformers.components.attention import Attention, AttentionConfig, register_attention
   
  @@ -22,7 +23,9 @@
           self.requires_input_projection = False
   
       def forward(self, q: torch.Tensor, *_, **__):
  -        att = torch.fft.fft2(q).real
  +        # Guard against autocast / fp16, not supported by torch.fft.fft2
  +        with autocast(enabled=False):
  +            att = torch.fft.fft2(q).real
   
           att = self.attn_drop(att)
  
verification_info: {"golden_diff": "diff --git a/xformers/components/attention/fourier_mix.py b/xformers/components/attention/fourier_mix.py\n--- a/xformers/components/attention/fourier_mix.py\n+++ b/xformers/components/attention/fourier_mix.py\n@@ -4,6 +4,7 @@\n # LICENSE file in the root directory of this source tree.\n \n import torch\n+from torch.cuda.amp import autocast\n \n from xformers.components.attention import Attention, AttentionConfig, register_attention\n \n@@ -22,7 +23,9 @@\n         self.requires_input_projection = False\n \n     def forward(self, q: torch.Tensor, *_, **__):\n-        att = torch.fft.fft2(q).real\n+        # Guard against autocast / fp16, not supported by torch.fft.fft2\n+        with autocast(enabled=False):\n+            att = torch.fft.fft2(q).real\n \n         att = self.attn_drop(att)\n", "issue": "fourier_mix attention does not support mixed precision training\nFourier mix attention does not support half precision.\r\n\r\n```Traceback (most recent call last):\r\n ...\r\n  File \"/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\", line 1102, in _call_impl\r\n    return forward_call(*input, **kwargs)\r\n  File \"/content/muzik/src/modules/xformer/custom_block_factory.py\", line 196, in forward\r\n    x = self.mha(q, k, v, att_mask=att_mask)\r\n  File \"/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\", line 1102, in _call_impl\r\n    return forward_call(*input, **kwargs)\r\n  File \"/content/muzik/src/modules/xformer/custom_multi_head_dispatch.py\", line 254, in forward\r\n    y = self.attention(q=q, k=k, v=v, att_mask=att_mask)\r\n  File \"/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\", line 1102, in _call_impl\r\n    return forward_call(*input, **kwargs)\r\n  File \"/usr/local/lib/python3.7/dist-packages/xformers/components/attention/fourier_mix.py\", line 25, in forward\r\n    att = torch.fft.fft2(q).real\r\nRuntimeError: Unsupported dtype Half```\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\n\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\n\n\n@register_attention(\"fourier_mix\", AttentionConfig)\nclass FourierMix(Attention):\n    def __init__(self, dropout: float, *_, **__):\n        \"\"\"\n        FFT-based pseudo-attention mechanism, from\n        \"\n        \"FNet: Mixing Tokens with Fourier Transforms\"\n        Lee-Thorp et al., 2021, https://arxiv.org/pdf/2105.03824.pdf\n        \"\"\"\n        super().__init__()\n        self.attn_drop = torch.nn.Dropout(dropout, inplace=False)\n        self.requires_input_projection = False\n\n    def forward(self, q: torch.Tensor, *_, **__):\n        att = torch.fft.fft2(q).real\n\n        att = self.attn_drop(att)\n\n        return att\n", "path": "xformers/components/attention/fourier_mix.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nfrom torch.cuda.amp import autocast\n\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\n\n\n@register_attention(\"fourier_mix\", AttentionConfig)\nclass FourierMix(Attention):\n    def __init__(self, dropout: float, *_, **__):\n        \"\"\"\n        FFT-based pseudo-attention mechanism, from\n        \"\n        \"FNet: Mixing Tokens with Fourier Transforms\"\n        Lee-Thorp et al., 2021, https://arxiv.org/pdf/2105.03824.pdf\n        \"\"\"\n        super().__init__()\n        self.attn_drop = torch.nn.Dropout(dropout, inplace=False)\n        self.requires_input_projection = False\n\n    def forward(self, q: torch.Tensor, *_, **__):\n        # Guard against autocast / fp16, not supported by torch.fft.fft2\n        with autocast(enabled=False):\n            att = torch.fft.fft2(q).real\n\n        att = self.attn_drop(att)\n\n        return att\n", "path": "xformers/components/attention/fourier_mix.py"}]}