Add fallback for Protocol import

2023-01-07 13:08:21 -07:00 · 2023-01-07 13:08:21 -07:00 · 984b86dd0a
commit 984b86dd0a
parent 151233399c
1 changed files with 7 additions and 1 deletions
--- 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,