Update x_transformer.py

x_transformer.py

@@ -27,6 +27,7 @@ from functools import partial
 import torch
 from torch import nn, einsum, Tensor
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from collections import namedtuple
 from functools import wraps
@@ -206,7 +207,7 @@ class Attend(nn.Module):
 
         # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale
         
-        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+        with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
             out = F.scaled_dot_product_attention(
                 q, k, v,
                 attn_mask = mask,