Type annotation correctness for create_act
parent
7d121ac2ef
commit
7bc7798d0e
|
@ -148,7 +148,7 @@ def get_act_layer(name: Union[Type[nn.Module], str] = 'relu'):
|
|||
return _ACT_LAYER_DEFAULT[name]
|
||||
|
||||
|
||||
def create_act_layer(name: Union[nn.Module, str], inplace=None, **kwargs):
|
||||
def create_act_layer(name: Union[Type[nn.Module], str], inplace=None, **kwargs):
|
||||
act_layer = get_act_layer(name)
|
||||
if act_layer is None:
|
||||
return None
|
||||
|
|
Loading…
Reference in New Issue