Change order of checkpoitn filtering fn application in builder, try dict, model variant first

pull/1641/head
Ross Wightman 2 years ago
parent e9f1376cde
commit 32f252381d

@ -179,11 +179,11 @@ def load_pretrained(
return return
if filter_fn is not None: if filter_fn is not None:
# for backwards compat with filter fn that take one arg, try one first, the two
try: try:
state_dict = filter_fn(state_dict)
except TypeError:
state_dict = filter_fn(state_dict, model) state_dict = filter_fn(state_dict, model)
except TypeError as e:
# for backwards compat with filter fn that take one arg
state_dict = filter_fn(state_dict)
input_convs = pretrained_cfg.get('first_conv', None) input_convs = pretrained_cfg.get('first_conv', None)
if input_convs is not None and in_chans != 3: if input_convs is not None and in_chans != 3:

Loading…
Cancel
Save