import torch, os
ckpt = torch.load('/home/ubuntu/sooktam2/model_1250000.pt', map_location='cpu', weights_only=False)
print(f'Before: update={ckpt["update"]}')
ckpt['update'] = 0
torch.save(ckpt, '/home/ubuntu/sooktam2_reset.pt')
v = torch.load('/home/ubuntu/sooktam2_reset.pt', map_location='cpu', weights_only=False)
print(f'After: update={v["update"]}')
print(f'Size: {os.path.getsize("/home/ubuntu/sooktam2_reset.pt")/1e9:.2f}GB')
print('DONE')