onnx

游雁

2023-02-25 9ccbadc8be2b52add807c421aa766a94e9176e44

 funasr/export/export_model.py

@@ -58,7 +58,7 @@
        if enc_size:
            dummy_input = model.get_dummy_inputs(enc_size)
        else:
            dummy_input = model.get_dummy_inputs_txt()
            dummy_input = model.get_dummy_inputs()

        # model_script = torch.jit.script(model)
        model_script = torch.jit.trace(model, dummy_input)
@@ -111,7 +111,7 @@
            dummy_input,
            os.path.join(path, f'{model.model_name}.onnx'),
            verbose=verbose,
            opset_version=12,
            opset_version=14,
            input_names=model.get_input_names(),
            output_names=model.get_output_names(),
            dynamic_axes=model.get_dynamic_axes()