Skip to content

Commit

Permalink
Bugfix: Can't serialize models without onnx_data file
Browse files Browse the repository at this point in the history
- 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
  • Loading branch information
DevinTDHa committed Mar 8, 2024
1 parent fa41371 commit b73dc0b
Showing 1 changed file with 5 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
}
}

Expand Down

0 comments on commit b73dc0b

Please sign in to comment.