Skip to content

Commit

Permalink
final_epoch EMA bug fix (#2317)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored Feb 27, 2021
1 parent cbd55da commit dfeec19
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion train.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down

0 comments on commit dfeec19

Please sign in to comment.