final_epoch EMA bug fix (ultralytics#2317)

KMint1819 · Feb 27, 2021 · 6ef0076 · 6ef0076
1 parent 7a6bba7
commit 6ef0076
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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}