From 1cfd37045572a575268bbb07f2fbbc871d05c884 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 22 Jul 2019 02:18:34 +0200
Subject: [PATCH] updates

---
 train.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index a045241f..c455d868 100644
--- a/train.py
+++ b/train.py
@@ -235,9 +235,10 @@ def train(cfg,
 
             # SGD burn-in
             if epoch == 0 and i <= n_burnin:
-                lr = hyp['lr0'] * (i / n_burnin) ** 4
+                g = (i / n_burnin) ** 4  # gain
                 for x in optimizer.param_groups:
-                    x['lr'] = lr
+                    x['lr'] = hyp['lr0'] * g
+                    x['weight_decay'] = hyp['weight_decay'] * g
 
             # Run model
             pred = model(imgs)