Model weight preservation in the Train diffused estimator (SwinIR) stage.

[24wenjie-li]

I retrained the Train diffused estimator (SwinIR) stage and found that the saved weights are dict_keys(['iter_start','log_step','log_step_img','state_dict']) and not SwinIR's weights, so the inference reports a load model error. Would like to ask if the code in main_sr.py about saving model weights is wrong? I'm looking forward to hearing from the author!

[zsyOAOA]

You should load the key of "state_dict". I have extracted this key for the released model weight.

[24wenjie-li]

You should load the key of "state_dict". I have extracted this key for the released model weight.

Thanks for the tip!