I am working on a multiclass segmentation problem using a Unet based approach. As loss function I use cross entropy, but for validation purposes dice and IoU are calculated too. Currently, the weights are stored (and overwritten) after each epoch. Therefore, I end up with the weights of the last epoch, which are not necessarily the best. Now I want to implement that the best model is saved.
My question now is, on what value should the decision for the best model be based? The loss, which is used to train the model, or the dice which is the metric I want to maximize. Is there a universal answer to this, or does it differ from problem to problem?