File size: 588 Bytes
3c38b94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | """Run this once before uploading. Strips optimizer state and saves a clean
safetensors file + config.json for Hugging Face."""
import json, torch
from safetensors.torch import save_file
ck = torch.load('big.pt', map_location='cpu')
sd = ck['model']
# safetensors needs contiguous tensors and no shared storage (untie weights)
sd = {k: v.clone().contiguous() for k, v in sd.items()}
save_file(sd, 'model.safetensors')
json.dump(ck['cfg'], open('config.json', 'w'), indent=2)
print("wrote model.safetensors + config.json")
print("config:", ck['cfg'], "| trained to iter", ck.get('iter'))
|