'''
Created on Mar 22, 2020

@author: Michal.Busta at gmail.com
'''

#get rid of the optimizer state ... 
import torch
MODEL_PATH = '/models/model-b2-2.pth'
state = torch.load(MODEL_PATH, map_location=lambda storage, loc: storage)

state_out = {
  "state_dict": state["state_dict"],
}

torch.save(state_out, 'model-b2-2.pth')
[((10, 8, 10, 73), 'torch.load', 'torch.load', (), '', False, 'import torch\n'), ((16, 0, 16, 39), 'torch.save', 'torch.save', ({(16, 11, 16, 20): 'state_out', (16, 22, 16, 38): '"""model-b2-2.pth"""'}, {}), "(state_out, 'model-b2-2.pth')", False, 'import torch\n')]