iumyx2612 commited on
Commit
a42af30
·
unverified ·
1 Parent(s): 922fbd8

Update `strip_optimizer()` (#5949)

Browse files

Replace 'training_result' with 'best_fitness' in strip_optimizer() to match key with ckpt from train.py

Files changed (1) hide show
  1. utils/general.py +1 -1
utils/general.py CHANGED
@@ -738,7 +738,7 @@ def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_op
738
  x = torch.load(f, map_location=torch.device('cpu'))
739
  if x.get('ema'):
740
  x['model'] = x['ema'] # replace model with ema
741
- for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
742
  x[k] = None
743
  x['epoch'] = -1
744
  x['model'].half() # to FP16
 
738
  x = torch.load(f, map_location=torch.device('cpu'))
739
  if x.get('ema'):
740
  x['model'] = x['ema'] # replace model with ema
741
+ for k in 'optimizer', 'best_fitness', 'wandb_id', 'ema', 'updates': # keys
742
  x[k] = None
743
  x['epoch'] = -1
744
  x['model'].half() # to FP16