Train own dataset, load cINN model.'state_dict' is error

[bjz0926]

mod stage2_cINN/main.py file,
if PFVD < best_PFVD:
torch.save(save_dict, save_path + '/checkpoint_best_val.pth')

---

if PFVD < best_PFVD:
torch.save(network.state_dict(), save_path + '/checkpoint_best_val.pth')

by python -W ignore -m stage2_cINN.main -gpu <gpu_id> -cf stage2_cINN/configs/ I get the checkpoint_best_val.pth. I change its name to cINN.pth, but error. Please tell the main.py that is right or error.
Thank you