From 6ef007658afef345752b303232a750a39114f0fa Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 27 Feb 2021 12:51:33 -0800 Subject: [PATCH] final_epoch EMA bug fix (#2317) --- train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index bbf879f3af5f..5c203f12651d 100644 --- a/train.py +++ b/train.py @@ -383,7 +383,8 @@ def train(hyp, opt, device, tb_writer=None, wandb=None): ckpt = {'epoch': epoch, 'best_fitness': best_fitness, 'training_results': results_file.read_text(), - 'model': deepcopy(model.module if is_parallel(model) else model).half(), + 'model': ema.ema if final_epoch else deepcopy( + model.module if is_parallel(model) else model).half(), 'ema': (deepcopy(ema.ema).half(), ema.updates), 'optimizer': optimizer.state_dict(), 'wandb_id': wandb_run.id if wandb else None}