Make xlformers optional

pytorch · Jun 12, 2024 · df9f3d9 · df9f3d9
1 parent e915cf9
commit df9f3d9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchbenchmark/operators/flash_attention/operator.py b/torchbenchmark/operators/flash_attention/operator.py
@@ -208,7 +208,7 @@ def xformers_preprocess(
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
-    ) -> xformers_fmha.Inputs:
+    ):
         q_1, k_1, v_1 = permute_qkv(q, k, v, perm=(0, 2, 1, 3))
         attn_bias = xformers.ops.LowerTriangularMask() if self.causal else None
         fhma_input = xformers_fmha.Inputs(