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,
|
ckpt = {'epoch': epoch,
|
||||||
'best_fitness': best_fitness,
|
'best_fitness': best_fitness,
|
||||||
'training_results': results_file.read_text(),
|
'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),
|
'ema': (deepcopy(ema.ema).half(), ema.updates),
|
||||||
'optimizer': optimizer.state_dict(),
|
'optimizer': optimizer.state_dict(),
|
||||||
'wandb_id': wandb_run.id if wandb else None}
|
'wandb_id': wandb_run.id if wandb else None}
|
||||||
|
|||||||
Reference in new issue
Block a user