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'))