Add fallback for Protocol import

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,