glenn-jocher commited on
Commit
1014583
1 Parent(s): cce95e7

strip_optimizer() default to save FP16 model

Browse files
Files changed (1) hide show
  1. utils/utils.py +1 -0
utils/utils.py CHANGED
@@ -620,6 +620,7 @@ def strip_optimizer(f='weights/best.pt'): # from utils.utils import *; strip_op
620
  # Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
621
  x = torch.load(f, map_location=torch.device('cpu'))
622
  x['optimizer'] = None
 
623
  torch.save(x, f)
624
  print('Optimizer stripped from %s' % f)
625
 
 
620
  # Strip optimizer from *.pt files for lighter files (reduced by 1/2 size)
621
  x = torch.load(f, map_location=torch.device('cpu'))
622
  x['optimizer'] = None
623
+ x['model'].half() # to FP16
624
  torch.save(x, f)
625
  print('Optimizer stripped from %s' % f)
626