glenn-jocher
commited on
Commit
•
dfeec19
1
Parent(s):
cbd55da
final_epoch EMA bug fix (#2317)
Browse files
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':
|
|
|
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}
|