From b73dc0b1ecdb49af9f2fa6e47b0af23d47442a53 Mon Sep 17 00:00:00 2001 From: Devin Ha Date: Fri, 8 Mar 2024 18:12:09 +0100 Subject: [PATCH] Bugfix: Can't serialize models without onnx_data file - onnxModelPath is not set for models without an .onnx_data file, so it will be None - None.get will throw an error, this checks for it first --- .../com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala index 2735626930016d..5c9156539d4cd0 100644 --- a/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala +++ b/src/main/scala/com/johnsnowlabs/ml/onnx/OnnxSerializeModel.scala @@ -54,10 +54,11 @@ trait WriteOnnxModel { fs.copyFromLocalFile(new Path(onnxFile), new Path(path)) // 4. check if there is a onnx_data file - - val onnxDataFile = new Path(onnxWrapper.onnxModelPath.get + dataFileSuffix) - if (fs.exists(onnxDataFile)) { - fs.copyFromLocalFile(onnxDataFile, new Path(path)) + if (onnxWrapper.onnxModelPath.isDefined) { + val onnxDataFile = new Path(onnxWrapper.onnxModelPath.get + dataFileSuffix) + if (fs.exists(onnxDataFile)) { + fs.copyFromLocalFile(onnxDataFile, new Path(path)) + } } }