In the test function in train_fns.py, when saving the best model based on best IS or FID, the best model is first saved, and then the state_dict is updated with the new best FID and IS values. Therefore, in the saved model's state_dict, the best FID and IS values of the previous best model is shown, instead of the current model.
This pull request will solve the issue by updating the best metrics in the state_dict before saving the best model.
In the test function in train_fns.py, when saving the best model based on best IS or FID, the best model is first saved, and then the state_dict is updated with the new best FID and IS values. Therefore, in the saved model's state_dict, the best FID and IS values of the previous best model is shown, instead of the current model. This pull request will solve the issue by updating the best metrics in the state_dict before saving the best model.