From 11b13e94a3d9ec1532816701e7f554f06bf217c1 Mon Sep 17 00:00:00 2001 From: Julien Chaumond Date: Fri, 24 Jan 2020 14:00:57 -0500 Subject: [PATCH] Add type to help my IDE out --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a9074d749c..7ea606b256 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -512,7 +512,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin): # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants # so we need to apply the function recursively. - def load(module, prefix=""): + def load(module: nn.Module, prefix=""): local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) module._load_from_state_dict( state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs