[Fix] use Callable instead of type[xxx] (#65700)

megemini · web-flow · commit d8ce75606ee4 · 2024-07-04T18:23:00.000+08:00
diff --git a/python/paddle/vision/models/mobilenetv2.py b/python/paddle/vision/models/mobilenetv2.py
@@ -16,6 +16,7 @@
 
 from typing import (
     TYPE_CHECKING,
+    Callable,
     TypedDict,
 )
 
@@ -53,7 +54,7 @@ def __init__(
         oup: int,
         stride: int,
         expand_ratio: float,
-        norm_layer: type[nn.Layer] = nn.BatchNorm2D,
+        norm_layer: Callable[..., nn.Layer] = nn.BatchNorm2D,
     ) -> None:
         super().__init__()
         self.stride = stride
diff --git a/python/paddle/vision/models/mobilenetv3.py b/python/paddle/vision/models/mobilenetv3.py
@@ -17,6 +17,7 @@
 from functools import partial
 from typing import (
     TYPE_CHECKING,
+    Callable,
     TypedDict,
 )
 
@@ -69,8 +70,8 @@ def __init__(
         self,
         input_channels: int,
         squeeze_channels: int,
-        activation: type[nn.Layer] = nn.ReLU,
-        scale_activation: type[nn.Layer] = nn.Sigmoid,
+        activation: Callable[..., nn.Layer] = nn.ReLU,
+        scale_activation: Callable[..., nn.Layer] = nn.Sigmoid,
     ) -> None:
         super().__init__()
         self.avgpool = nn.AdaptiveAvgPool2D(1)
@@ -136,8 +137,8 @@ def __init__(
         filter_size: int,
         stride: int,
         use_se: bool,
-        activation_layer: type[nn.Layer],
-        norm_layer: type[nn.Layer],
+        activation_layer: Callable[..., nn.Layer],
+        norm_layer: Callable[..., nn.Layer],
     ) -> None:
         super().__init__()
         self.use_res_connect = stride == 1 and in_channels == out_channels
diff --git a/python/paddle/vision/models/resnet.py b/python/paddle/vision/models/resnet.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import paddle
 from paddle import nn
@@ -121,7 +121,7 @@ def __init__(
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: type[nn.Layer] | None = None,
+        norm_layer: Callable[..., nn.Layer] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -173,7 +173,7 @@ def __init__(
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: type[nn.Layer] | None = None,
+        norm_layer: Callable[..., nn.Layer] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
diff --git a/python/paddle/vision/models/shufflenetv2.py b/python/paddle/vision/models/shufflenetv2.py
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import paddle
 from paddle import nn
@@ -121,7 +121,7 @@ def __init__(
         in_channels: int,
         out_channels: int,
         stride: Size2,
-        activation_layer: type[nn.Layer] = nn.ReLU,
+        activation_layer: Callable[..., nn.Layer] = nn.ReLU,
     ) -> None:
         super().__init__()
         self._conv_pw = ConvNormActivation(
@@ -171,7 +171,7 @@ def __init__(
         in_channels: int,
         out_channels: int,
         stride: Size2,
-        activation_layer: type[nn.Layer] = nn.ReLU,
+        activation_layer: Callable[..., nn.Layer] = nn.ReLU,
     ) -> None:
         super().__init__()
 
diff --git a/python/paddle/vision/ops.py b/python/paddle/vision/ops.py
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal, Sequence, overload
+from typing import TYPE_CHECKING, Callable, Literal, Sequence, overload
 
 import numpy as np
 
@@ -1902,8 +1902,8 @@ def __init__(
         stride: Size2 = 1,
         padding: _PaddingSizeMode | Size2 | Size4 | str | None = None,
         groups: int = 1,
-        norm_layer: type[nn.Layer] = BatchNorm2D,
-        activation_layer: type[nn.Layer] = ReLU,
+        norm_layer: Callable[..., nn.Layer] = BatchNorm2D,
+        activation_layer: Callable[..., nn.Layer] = ReLU,
         dilation: int = 1,
         bias: bool | None = None,
     ) -> None:
