pytorch2torchscript.py

from model_ecapatdnn import ECAPAModel
import soundfile as sf
import torch

# load your pytorch model (+ weights) here
model_1 = ECAPAModel.ECAPAModel(lr = 0.001, lr_decay = 0.97, C = 1024, n_class = 18505, m = 0.2, s = 30, test_step = 3, gpu = -1)
model_1.load_parameters("/ecapatdnn/exps/finetuned/model/model_0028.model")

model = model_1.speaker_encoder # module with forward implemented
# if your forward function contains any arguments other than the input tensors, make sure to add default values
# Switch the model to eval model
model.eval()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 48000)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

# Save the TorchScript model
traced_script_module.save("model.pt")