aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--modules/sub_quadratic_attention.py8
1 files changed, 7 insertions, 1 deletions
diff --git a/modules/sub_quadratic_attention.py b/modules/sub_quadratic_attention.py
index fea7aaac..93381bae 100644
--- a/modules/sub_quadratic_attention.py
+++ b/modules/sub_quadratic_attention.py
@@ -15,7 +15,13 @@ import torch
from torch import Tensor
from torch.utils.checkpoint import checkpoint
import math
-from typing import Optional, NamedTuple, Protocol, List
+
+try:
+ from typing import Protocol
+except:
+ from typing_extensions import Protocol
+
+from typing import Optional, NamedTuple, List
def narrow_trunc(
input: Tensor,