fix format

facebookresearch · May 3, 2024 · 50d747c · 50d747c
1 parent 2ea2b49
commit 50d747c
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/xformers/benchmarks/benchmark_attn_decoding.py b/xformers/benchmarks/benchmark_attn_decoding.py
@@ -180,8 +180,17 @@ class AttentionDecodingCKSplitKV(AttentionDecodingBase):
 
 class AttentionDecodingSplitInt4KV(AttentionDecodingBase):
     OP = xops.fmha.triton_splitk.FwOp
-    def __init__(self, B: int, Mq: int, Mkv: int, Hq: int, Hkv: int, K: int, bw: bool,
-                 attn_bias_type
+
+    def __init__(
+        self,
+        B: int,
+        Mq: int,
+        Mkv: int,
+        Hq: int,
+        Hkv: int,
+        K: int,
+        bw: bool,
+        attn_bias_type,
     ) -> None:
         # super(AttentionDecodingSplitInt4KV, self).__init__(B, Mq, Mkv, Hq, Hkv, K, bw, attn_bias_type)
         dtype = torch.float16
@@ -255,7 +264,7 @@ def __init__(self, B: int, Mq: int, Mkv: int, Hq: int, Hkv: int, K: int, bw: boo
             not_supported_reasons = self.OP.not_supported_reasons(inp)
             if not_supported_reasons:
                 raise NotSupportedInputError(not_supported_reasons)
-            
+
 
 class AttentionDecodingPyTorchRepeat(AttentionDecodingBase):
     def fw(self) -> None: