| | |
| | | kwargs["device"] = device |
| | | |
| | | torch.set_num_threads(kwargs.get("ncpu", 4)) |
| | | |
| | | |
| | | # build tokenizer |
| | | tokenizer = kwargs.get("tokenizer", None) |
| | |
| | | export_dir = export_utils.export_onnx( |
| | | model=model, |
| | | data_in=data_list, |
| | | quantize=quantize, |
| | | fallback_num=fallback_num, |
| | | calib_num=calib_num, |
| | | opset_version=opset_version, |
| | | **kwargs) |
| | | else: |
| | | export_dir = export_utils.export_torchscripts( |
| | | model=model, |
| | | data_in=data_list, |
| | | quantize=quantize, |
| | | fallback_num=fallback_num, |
| | | calib_num=calib_num, |
| | | opset_version=opset_version, |
| | | **kwargs) |
| | | |
| | | return export_dir |