infer.py doesn't work, variable referenced before assignment

[PT0X0E]

Hi, thanks for the work!
But I think there is an obvious bug at line 56 in src_files/models/utils/factory.py:
model.load_state_dict(state[key], strict=True)
The variable key in else branch is used without definition. So when load_head argument is True, which True in infer.py, the code doesn't work.

[BigBuffa1o]

Move the if else loop out of the load head loop will solve this problem.

        if 'model' in state:
            key = 'model'
        else:
            key = 'state_dict'
        if not load_head:
            filtered_dict = {k: v for k, v in state[key].items() if
                             (k in model.state_dict() and 'head.fc' not in k)}
            model.load_state_dict(filtered_dict, strict=False)
        else:
            model.load_state_dict(state[key], strict=True)

[yangyangtiaoguo]

Move the if else loop out of the load head loop will solve this problem.

        if 'model' in state:
            key = 'model'
        else:
            key = 'state_dict'
        if not load_head:
            filtered_dict = {k: v for k, v in state[key].items() if
                             (k in model.state_dict() and 'head.fc' not in k)}
            model.load_state_dict(filtered_dict, strict=False)
        else:
            model.load_state_dict(state[key], strict=True)

Still does not work, model loading failed.
RuntimeError: Error(s) in loading state_dict for TResNet:

[KoSweet]

The solution works for me. Thanks