Files
audio_model/src/handle/export_model_state_dict.py
2026-05-07 11:29:21 +06:00

13 lines
398 B
Python

from pathlib import Path
import torch
device = "cuda:0"
workspace_dir = Path(__file__).parent.parent.parent
input_checkpoint = workspace_dir.joinpath('.checkpoints/best_wer_model.pt')
output_checkpoint = workspace_dir.joinpath('.checkpoints/prodect_best_wer_model.pt')
checkpoint = torch.load(input_checkpoint, map_location=device)
torch.save(checkpoint['model_state_dict'], output_checkpoint)