From c1c09eb3ccf18c74ddf0c4217ab306e0a60c0775 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 10 Mar 2019 15:03:17 +0100
Subject: [PATCH] Update train.py

---
 train.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/train.py b/train.py
index f53ddf60..6b7b6181 100644
--- a/train.py
+++ b/train.py
@@ -151,9 +151,8 @@ def train(
             print(s)
 
         # Update best loss
-        loss_per_target = rloss['loss'] / rloss['nT']
-        if loss_per_target < best_loss:
-            best_loss = loss_per_target
+        if rloss['loss'] < best_loss:
+            best_loss = rloss['loss']
 
         # Save latest checkpoint
         checkpoint = {'epoch': epoch,
@@ -163,7 +162,7 @@ def train(
         torch.save(checkpoint, latest)
 
         # Save best checkpoint
-        if best_loss == loss_per_target:
+        if best_loss == rloss['loss']:
             os.system('cp ' + latest + ' ' + best)
 
         # Save backup weights every 5 epochs (optional)