Error in prediction step after imported a keras model to tensorflow java

my goal is to use a Keras model in a java program. I export the keras model with model.export() and not model.save() so I get well a folder with the model in .pb format.

Then I used py .\saved_model_cli.py show -- dir '.' -all to see the inputs and outputs to fill in the java code. I get that :

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['__saved_model_init_op']:
  The given SavedModel SignatureDef contains the following input(s):
  The given SavedModel SignatureDef contains the following output(s):
    outputs['__saved_model_init_op'] tensor_info:
        dtype: DT_INVALID
        shape: unknown_rank
        name: NoOp
  Method name is:

signature_def['serve']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['keras_tensor'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 6)
        name: serve_keras_tensor:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['keras_tensor'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 6)
        name: serving_default_keras_tensor:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall_1:0
  Method name is: tensorflow/serving/predict
The MetaGraph with tag set ['serve'] contains the following ops: {'ReadVariableOp', 'Select', 'StatefulPartitionedCall', 'RestoreV2', 'NoOp', 'Identity', 'StaticRegexFullMatch', 'StringJoin', 'AssignVariableOp', 'SaveV2', 'MergeV2Checkpoints', 'VarIsInitializedOp', 'AddV2', 'VarHandleOp', 'DisableCopyOnRead', 'Pack', 'Placeholder', 'MatMul', 'Const', 'Relu', 'ShardedFilename'}

Concrete Functions:2024-11-12 16:47:24.597134: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.

  Function Name: 'serve'
    Option #1
      Callable with:
        Argument #1
          keras_tensor: TensorSpec(shape=(None, 6), dtype=tf.float32, name='keras_tensor')

Finally, the java code to import and make a prediction is :

public static void importKerasModel() {
        try (SavedModelBundle model = SavedModelBundle.load("PATH\kerasModel", "serve")) {
            float[] x = {0.48f, 0.48f, 0.48f, 0.48f, 0.48f, 0.48f};
            try (Tensor input = TFloat32.vectorOf(x);
                 Tensor output = model.session()
                         .runner()
                         .feed("serve_keras_tensor", input)
                         .fetch("StatefulPartitionedCall")
                         .run()
                         .get(0)) {

                float prediction = output.dataType().getNumber();
                System.out.println("prediction = " + prediction);
            }
        }
    }

But I get this error message :

2024-11-12 17:26:01.089591: I tensorflow/cc/saved_model/loader.cc:317] SavedModel load for tags { serve }; Status: success: OK. Took 61548 microseconds.
2024-11-12 17:26:01.317247: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: In[0] is not a matrix
     [[{{node StatefulPartitionedCall/StatefulPartitionedCall/sequential_1/dense_1/Relu}}]]
Exception in thread "main" org.tensorflow.exceptions.TFInvalidArgumentException: In[0] is not a matrix
     [[{{node StatefulPartitionedCall/StatefulPartitionedCall/sequential_1/dense_1/Relu}}]]
    at org.tensorflow.internal.c_api.AbstractTF_Status.throwExceptionIfNotOK(AbstractTF_Status.java:76)
    at org.tensorflow.Session.run(Session.java:826)
    at org.tensorflow.Session$Runner.runHelper(Session.java:549)
    at org.tensorflow.Session$Runner.run(Session.java:476)
    at com.ptvgroup.platform.truckslogs.converter.HelloTensorFlow.importKerasModel(HelloTensorFlow.java:471)
    at com.ptvgroup.platform.truckslogs.converter.Main.main(Main.java:25)

Someone can help me ? What does it means "In[0] is not a matrix" ? It's because my dimensions/shapes of the inputs and outputs are (-1,6) and (-1,1) ?

The error 'In[0] is not a matrix' comes from a wrong data type when I would like make a prediction. I create a tensor of vector instead a tensor of matrix. The exception message told well the value was not a matrix. The value concerned is the tensor of inputs.

public static void importKerasModel() {
        try (SavedModelBundle model = SavedModelBundle.load("PATH", "serve")) {
            float[] x = {200f,0f,1.5f,2f,2.5f,0f};
            FloatNdArray matrix = NdArrays.ofFloats(Shape.of(1, 6)); // my model have 6 features per observations
            matrix.set(NdArrays.vectorOf(x), 0);
            try (Tensor input = TFloat32.tensorOf(matrix);
                 TFloat32 output = (TFloat32) model.session()
                         .runner()
                         .feed("serve_keras_tensor_272", input)  ///## to know inputs and outputs  py .\saved_model_cli.py show --dir '.' --all
                         .fetch("StatefulPartitionedCall")
                         .run()
                         .get(0)) {
                float prediction =  output.getFloat();;
                System.out.println("prediction = " + prediction);
            }
        }
    }

This code fix the failed execution.