mlp_bias Default = True

flowvision/layers/attention/se.py

@@ -1,6 +1,8 @@
+from typing import Optional
+
 import oneflow as flow
 import oneflow.nn as nn
-
+from oneflow.nn import ReLU, Sigmoid
 
 class SEModule(nn.Module):
     """
@@ -11,19 +13,19 @@ class SEModule(nn.Module):
         channels (int): The input channel size
         reduction (int): Ratio that allows us to vary the capacity and computational cost of the SE Module. Default: 16
         rd_channels (int or None): Number of reduced channels. If none, uses reduction to calculate
-        act_layer (flow.nn.Module): An activation layer used after the first FC layer. Default: flow.nn.ReLU
-        gate_layer (flow.nn.Module): An activation layer used after the second FC layer. Default: flow.nn.Sigmoid
-        mlp_bias (bool): If True, add learnable bias to the linear layers. Default: False
+        act_layer (Optional[ReLU]): An activation layer used after the first FC layer. Default: flow.nn.ReLU
+        gate_layer (Optional[Sigmoid]): An activation layer used after the second FC layer. Default: flow.nn.Sigmoid
+        mlp_bias (bool): If True, add learnable bias to the linear layers. Default: True
     """
 
     def __init__(
         self,
-        channels,
-        reduction=16,
-        rd_channels=None,
-        act_layer=nn.ReLU,
-        gate_layer=nn.Sigmoid,
-        mlp_bias=False,
+        channels: int,
+        reduction: int = 16,
+        rd_channels: int = None,
+        act_layer: Optional[ReLU] = nn.ReLU,
+        gate_layer: Optional[Sigmoid] = nn.Sigmoid,
+        mlp_bias=True,
     ):
         super(SEModule, self).__init__()
         rd_channels = channels // reduction if rd_channels is None else rd_channels