File tree Expand file tree Collapse file tree 1 file changed +11
-2
lines changed
backends/.jvm/src/main/scala Expand file tree Collapse file tree 1 file changed +11
-2
lines changed Original file line number Diff line number Diff line change @@ -116,8 +116,8 @@ trait ORTOperatorBackend extends OpToONNXBytesConverter with AutoCloseable {
116
116
Some (x.map { y =>
117
117
getOnnxTensor(y._1, y._2._3.toSeq.toArray, env)
118
118
})
119
+ case None => None
119
120
}
120
- case None => None
121
121
case tens : Tensor [T , Tuple3 [Tt , Td , S ]] =>
122
122
Some (tens.map { x =>
123
123
getOnnxTensor(x._1, x._2._3.toSeq.toArray, env)
@@ -155,7 +155,15 @@ trait ORTOperatorBackend extends OpToONNXBytesConverter with AutoCloseable {
155
155
} yield res(
156
156
opToModelProto(
157
157
opName,
158
- (t.map(_.getInfo.onnxType.value) zip t.map(_.getInfo.getShape.map(_.toInt))),
158
+ (t.map(_.getInfo.onnxType.value) zip { t.map(_.getInfo.getShape.map(_.toInt) match {
159
+ // ORT shape inference diverges from the ONNX spec in requiring a scalar here instead of a tensor with shape,
160
+ // causing a crash without this fix
161
+ case Array (1 ) => if (opName.equals(" Dropout" )) Array [Int ]() else Array (1 )
162
+ case y : Array [Int ] => y
163
+ }
164
+ )
165
+ }
166
+ ),
159
167
attrs
160
168
).toByteArray,
161
169
tens
@@ -185,6 +193,7 @@ trait ORTOperatorBackend extends OpToONNXBytesConverter with AutoCloseable {
185
193
opName,
186
194
attrs
187
195
)
196
+ // TODO: now that this is otherwise working, try memoizing here
188
197
result.flatMap(IO .println(" Real call opName => " + opName).as(_))
189
198
}
190
199
You can’t perform that action at this time.
0 commit comments