-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Closed
Description
Currently, _remap_standard_resnet_model is designed according to the network structure shown in https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html .
Therefore, we assume that there is not extra key (see also here ) to extract the expected state dict.
It would be nice to make the network be compatible with the weights that contain some other keys. For example, users may also save other information rather than state dict within the weights:
dict_keys(['epoch', 'arch', 'state_dict', 'optimizer', 'version', 'args', 'amp_scaler', 'metric'])
Therefore, I suggest we can add an arg to support this kind of case. In addition, we can also add some notifications within the load function to help users double check that if the load method is successful.
cc @KumoLiu
Metadata
Metadata
Assignees
Labels
No labels