Enforce IR version for older onnxruntime

This commit is contained in:
henryruhs
2025-03-11 14:43:09 +01:00
parent 3fe32d7832
commit f05ff6cdb1
6 changed files with 12 additions and 4 deletions
+5 -3
View File
@@ -13,10 +13,12 @@ def export() -> None:
directory_path = CONFIG.get('exporting', 'directory_path')
source_path = CONFIG.get('exporting', 'source_path')
target_path = CONFIG.get('exporting', 'target_path')
ir_version = CONFIG.getint('exporting', 'ir_version')
opset_version = CONFIG.getint('exporting', 'opset_version')
makedirs(directory_path, exist_ok = True)
embedding_converter_trainer = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu')
embedding_converter_trainer.eval()
model = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu')
model.eval()
model.ir_version = ir_version
input_tensor = torch.randn(1, 512)
torch.onnx.export(embedding_converter_trainer, input_tensor, target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = opset_version)
torch.onnx.export(model, input_tensor, target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = opset_version)