Add missing imports

Author: adamjstewart

timm/layers/attention2d.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Type, Union
 
 import torch
 from torch import nn as nn

timm/models/fastvit.py

@@ -7,7 +7,7 @@
 #
 import os
 from functools import partial
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Type, Union
 
 import torch
 import torch.nn as nn

timm/models/swin_transformer_v2.py

@@ -13,7 +13,7 @@
 # Written by Ze Liu
 # --------------------------------------------------------
 import math
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Type, Union
 
 import torch
 import torch.nn as nn

timm/models/vgg.py

@@ -5,7 +5,7 @@
 
 Copyright 2021 Ross Wightman
 """
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, Type, Union, cast
 
 import torch
 import torch.nn as nn