Skip to content

Commit 926dadf

Browse files
authored
final_epoch EMA bug fix (ultralytics#2317)
1 parent 2c6baf7 commit 926dadf

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

train.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -383,7 +383,8 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
383383
ckpt = {'epoch': epoch,
384384
'best_fitness': best_fitness,
385385
'training_results': results_file.read_text(),
386-
'model': deepcopy(model.module if is_parallel(model) else model).half(),
386+
'model': ema.ema if final_epoch else deepcopy(
387+
model.module if is_parallel(model) else model).half(),
387388
'ema': (deepcopy(ema.ema).half(), ema.updates),
388389
'optimizer': optimizer.state_dict(),
389390
'wandb_id': wandb_run.id if wandb else None}

0 commit comments

Comments
 (0)