save ckpt in FP16 #119

2020-06-17 22:34:13 -07:00 · 2020-06-17 22:34:13 -07:00 · d9b64c27c2
parent 9fdb0fbacf
commit d9b64c27c2
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -332,7 +332,7 @@ def train(hyp):
                ckpt = {'epoch': epoch,
                        'best_fitness': best_fitness,
                        'training_results': f.read(),
-                        'model': ema.ema.module if hasattr(model, 'module') else ema.ema,
+                        'model': ema.ema.module.half() if hasattr(model, 'module') else ema.ema.half(),
                        'optimizer': None if final_epoch else optimizer.state_dict()}

            # Save last, best and delete