final_epoch EMA bug fix (#2317)
This commit is contained in:
1 parent
cbd55da5d2
commit
dfeec198cb
1 file changed
+2
-1
@@ -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}
|
||||
|
||||
Reference in new issue
Block a user