Fix batch size detection for multi-input models (#572)

FremyCompany · caisq · commit a68baa535ca3 · 2019-08-03T09:47:45.000-04:00
diff --git a/src/engine/training.ts b/src/engine/training.ts
@@ -1130,7 +1130,7 @@ export class LayersModel extends Container implements tfc.InferenceModel {
     checkInputData(x, this.inputNames, this.feedInputShapes, true);
     // TODO(cais): Take care of the learning_phase boolean flag.
     //   if (this.useLearningPhase) ...
-    return this.predictLoop(x, x.shape[0]);
+    return this.predictLoop(x, Array.isArray(x) ? x[0].shape[0] : x.shape[0]);
   }
 
   protected standardizeUserDataXY(

Original file line number	Diff line number	Diff line change
`@@ -1130,7 +1130,7 @@ export class LayersModel extends Container implements tfc.InferenceModel {`
`1130`	`1130`	`checkInputData(x, this.inputNames, this.feedInputShapes, true);`
`1131`	`1131`	`// TODO(cais): Take care of the learning_phase boolean flag.`
`1132`	`1132`	`// if (this.useLearningPhase) ...`
`1133`		`- return this.predictLoop(x, x.shape[0]);`
	`1133`	`+ return this.predictLoop(x, Array.isArray(x) ? x[0].shape[0] : x.shape[0]);`
`1134`	`1134`	`}`
`1135`	`1135`
`1136`	`1136`	`protected standardizeUserDataXY(`