diff --git a/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/saved_model.pb b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/saved_model.pb new file mode 100644 index 00000000000..7f19c27d6a5 Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/saved_model.pb differ diff --git a/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..8020ed84834 Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 differ diff --git a/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.index b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.index new file mode 100644 index 00000000000..a2750a61168 Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/model_single_input_multi_output/variables/variables.index differ diff --git a/tfjs-node/src/saved_model.ts b/tfjs-node/src/saved_model.ts index 8df10f388a3..5b0bff7e157 100644 --- a/tfjs-node/src/saved_model.ts +++ b/tfjs-node/src/saved_model.ts @@ -284,9 +284,10 @@ export class TFSavedModel implements InferenceModel { let inputTensors: Tensor[] = []; if (inputs instanceof Tensor) { inputTensors.push(inputs); - return this.backend.runSavedModel( + const result = this.backend.runSavedModel( this.sessionId, inputTensors, Object.values(this.inputNodeNames), - Object.values(this.outputNodeNames))[0]; + Object.values(this.outputNodeNames)); + return result.length > 1 ? result : result[0]; } else if (Array.isArray(inputs)) { inputTensors = inputs; return this.backend.runSavedModel( diff --git a/tfjs-node/src/saved_model_test.ts b/tfjs-node/src/saved_model_test.ts index 3a9bed7a7b7..99fcaceddc5 100644 --- a/tfjs-node/src/saved_model_test.ts +++ b/tfjs-node/src/saved_model_test.ts @@ -429,7 +429,26 @@ describe('SavedModel', () => { model2.dispose(); }); - it('execute model with multiple inputs and outputs', async () => { + it('execute model with single input and multiple outputs', async () => { + // This test model behaves as: f(x)=[2*x, x] + const model = await tf.node.loadSavedModel( + './test_objects/saved_model/model_single_input_multi_output', ['serve'], + 'serving_default'); + const input = tf.tensor1d([1, 2, 3], 'int32'); + const output = model.predict(input) as tf.Tensor[]; + const output1 = output[0]; + const output2 = output[1]; + expect(output1.shape).toEqual(input.shape); + expect(output1.dtype).toBe(input.dtype); + expect(output2.shape).toEqual(input.shape); + expect(output2.dtype).toBe(input.dtype); + test_util.expectArraysClose(await output1.data(), [2, 4, 6]); + test_util.expectArraysClose(await output2.data(), [1, 2, 3]); + model.dispose(); + }); + + it('execute model with multiple inputs and multiple outputs', async () => { + // This test model behaves as: f(x, y)=[2*x, y] const model = await tf.node.loadSavedModel( './test_objects/saved_model/model_multi_output', ['serve'], 'serving_default'); diff --git a/tfjs-node/test_objects/saved_model/model_single_input_multi_output/saved_model.pb b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/saved_model.pb new file mode 100644 index 00000000000..7f19c27d6a5 Binary files /dev/null and b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/saved_model.pb differ diff --git a/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..8020ed84834 Binary files /dev/null and b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.data-00000-of-00001 differ diff --git a/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.index b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.index new file mode 100644 index 00000000000..a2750a61168 Binary files /dev/null and b/tfjs-node/test_objects/saved_model/model_single_input_multi_output/variables/variables.index differ