diff --git a/utils/utils.py b/utils/utils.py index 8f375b58..e4d67e4d 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -578,7 +578,8 @@ def print_model_biases(model): def strip_optimizer(f='weights/last.pt'): # from utils.utils import *; strip_optimizer() # Strip optimizer from *.pt files for lighter files (reduced by 2/3 size) - x = torch.load(f) + x = torch.load(f, map_location=torch.device('cpu')) + # x['epoch'] = 0 # uncomment to create a backbone x['optimizer'] = None torch.save(x, f)