EMA class updates

2020-03-14 17:34:13 -07:00 · 2020-03-14 17:34:13 -07:00 · 851c9b9883
commit 851c9b9883
parent d91469a516
1 changed files with 1 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -292,6 +292,7 @@ def train():
            if ni % accumulate == 0:
                optimizer.step()
                optimizer.zero_grad()
+                # ema.update(model)

            # Print batch results
            mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses