iumyx2612
commited on
Update `strip_optimizer()` (#5949)
Browse filesReplace 'training_result' with 'best_fitness' in strip_optimizer() to match key with ckpt from train.py
- 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', '
|
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
|