Update strip_optimizer() (#5949)

Replace 'training_result' with 'best_fitness' in strip_optimizer() to match key with ckpt from train.py
This commit is contained in:
iumyx2612 authored and GitHub committed 2021-12-10 18:06:27 +01:00
1 parent 922fbd8209
commit a42af30d8a
1 file changed
+1 -1
+1 -1
View File
@@ -738,7 +738,7 @@ def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_op
x = torch.load(f, map_location=torch.device('cpu'))
if x.get('ema'):
x['model'] = x['ema'] # replace model with ema
for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
for k in 'optimizer', 'best_fitness', 'wandb_id', 'ema', 'updates': # keys
x[k] = None
x['epoch'] = -1
x['model'].half() # to FP16