File size: 289 Bytes
0e08ab8
 
 
2b6c108
 
0e08ab8
 
 
 
 
2b6c108
1
2
3
4
5
6
7
8
9
10
11
12
import os

import torch
from safetensors.torch import save_file


checkpoint_dir = '../out/pretrain-core-3/hf'
output_dir = '../out/pretrain-core-3/hf'

state_dict = torch.load(os.path.join(checkpoint_dir, 'model.pth'))
save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))