use torch.jit.Final instead of Final for beit, eva

reinhud · May 12, 2023 · f744bda · f744bda
1 parent 35b9fc7
commit f744bda
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/timm/models/beit.py b/timm/models/beit.py
@@ -39,8 +39,7 @@
 # --------------------------------------------------------'
 
 import math
-from functools import partial
-from typing import Callable, Final, Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -80,7 +79,7 @@ def gen_relative_position_index(window_size: Tuple[int, int]) -> torch.Tensor:
 
 
 class Attention(nn.Module):
-    fused_attn: Final[bool]
+    fused_attn: torch.jit.Final[bool]
 
     def __init__(
             self,

diff --git a/timm/models/eva.py b/timm/models/eva.py
@@ -26,7 +26,7 @@
 # EVA02 models Copyright (c) 2023 BAAI-Vision
 
 import math
-from typing import Callable, Final, Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -45,7 +45,7 @@
 
 
 class EvaAttention(nn.Module):
-    fused_attn: Final[bool]
+    fused_attn: torch.jit.Final[bool]
 
     def __init__(
             self,