from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict


dir_path = "/home/minz/logs/ditto_v2_local/epoch=40-step=618362.ckpt"
out_path = "/home/minz/logs/ditto_v2_local/epoch=40.pt"

convert_zero_checkpoint_to_fp32_state_dict(dir_path, out_path)