Type annotation correctness for create_act

pull/2092/head
Ross Wightman 2024-02-10 14:57:58 -08:00
parent 7d121ac2ef
commit 7bc7798d0e
1 changed files with 1 additions and 1 deletions

View File

@ -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