Update seva/modules/transformer.py

seva/modules/transformer.py

@@ -67,8 +67,14 @@ class Attention(nn.Module):
             lambda t: rearrange(t, "b l (h d) -> b h l d", h=self.heads),
             (q, k, v),
         )
-        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+        with torch.backends.cuda.sdp_kernel(
+            enable_flash=False,
+            enable_mem_efficient=True,
+            enable_math=True,
+            enable_cudnn=True,
+        ):
             out = F.scaled_dot_product_attention(q, k, v)
+
         out = rearrange(out, "b h l d -> b l (h d)")
         out = self.to_out(out)
         return out