From 208b9a73fede99aef1d348db3dc25a7ed677bbf2 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 13 Mar 2020 16:08:49 -0700
Subject: [PATCH] updates

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

diff --git a/train.py b/train.py
index fe4d8a44..0a1d5525 100644
--- a/train.py
+++ b/train.py
@@ -240,12 +240,14 @@ def train():
             # Hyperparameter Burn-in
             n_burn = 200  # number of burn-in batches
             if ni <= n_burn:
-                # g = ni / n_burn  # gain
+                # g = (ni / n_burn) ** 2  # gain
                 for x in model.named_modules():  # initial stats may be poor, wait to track
                     if x[0].endswith('BatchNorm2d'):
                         x[1].track_running_stats = ni == n_burn
                 # for x in optimizer.param_groups:
                 #     x['lr'] = x['initial_lr'] * lf(epoch) * g  # gain rises from 0 - 1
+                #     if 'momentum' in x:
+                #         x['momentum'] = hyp['momentum'] * g
 
             # Plot images with bounding boxes
             if ni < 1: