[fix] FourierMix + AMP (#258)

Author: blefaudeux

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.0.x] - TBD
 ### Fixed
 - Fix some torchscriptability [#246]
+- Fix FourierMix being compatible with AMP [#258]
 
 ## [0.0.10] - 2022-03-14
 ### Fixed

tests/test_attentions.py

@@ -124,6 +124,11 @@ def test_order_invariance(
             att_2 = multi_head(inputs, inputs_shuffled, inputs)
             assert (att != att_2).any()
 
+        # Test AMP, if available
+        if device.type == "cuda":
+            with torch.cuda.amp.autocast(enabled=True):
+                _ = multi_head(inputs, inputs_shuffled, inputs)
+
 
 @pytest.mark.parametrize("heads", [1, 4])
 @pytest.mark.parametrize("attention_name", ["scaled_dot_product"])

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 @@ def __init__(self, dropout: float, *_, **__):
         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)