From d9b64c27c24db2001535bb480959aca015159510 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 17 Jun 2020 22:34:13 -0700 Subject: [PATCH] save ckpt in FP16 #119 --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 5163bb1..ce46797 100644 --- 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