Save best model based on train loss

2026-09-09 11:16:00 +08:00 · 2018-07-27 13:51:13 +02:00
parent 13ace3a9e7
commit 30561d0b5f
2 changed files with 4 additions and 3 deletions
@@ -410,7 +410,7 @@ def main(args):
        train_loss, current_step = train(model, criterion, criterion_st, train_loader, optimizer, optimizer_st, ap, epoch)
        val_loss = evaluate(model, criterion, criterion_st, val_loader, ap, current_step)
        print(" | > Train Loss: {:.5f}   Validation Loss: {:.5f}".format(train_loss, val_loss))
-        best_loss = save_best_model(model, optimizer, val_loss,
+        best_loss = save_best_model(model, optimizer, train_loss,
                                    best_loss, OUT_PATH,
                                    current_step, epoch)