glenn-jocher commited on
Commit
dfeec19
1 Parent(s): cbd55da

final_epoch EMA bug fix (#2317)

Browse files
Files changed (1) hide show
  1. train.py +2 -1
train.py CHANGED
@@ -383,7 +383,8 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
383
  ckpt = {'epoch': epoch,
384
  'best_fitness': best_fitness,
385
  'training_results': results_file.read_text(),
386
- 'model': deepcopy(model.module if is_parallel(model) else model).half(),
 
387
  'ema': (deepcopy(ema.ema).half(), ema.updates),
388
  'optimizer': optimizer.state_dict(),
389
  'wandb_id': wandb_run.id if wandb else None}
 
383
  ckpt = {'epoch': epoch,
384
  'best_fitness': best_fitness,
385
  'training_results': results_file.read_text(),
386
+ 'model': ema.ema if final_epoch else deepcopy(
387
+ model.module if is_parallel(model) else model).half(),
388
  'ema': (deepcopy(ema.ema).half(), ema.updates),
389
  'optimizer': optimizer.state_dict(),
390
  'wandb_id': wandb_run.id if wandb else None}