diff --git a/paddle/capi/examples/model_inference/dense/main.c b/paddle/capi/examples/model_inference/dense/main.c
index 3e6bd5285058a297c4574631e2a5c033b83936e8..876af2aa7615c098d225b56ce2ea0b1529a6e3c6 100644
--- a/paddle/capi/examples/model_inference/dense/main.c
+++ b/paddle/capi/examples/model_inference/dense/main.c
@@ -27,18 +27,20 @@ int main() {
   CHECK(paddle_arguments_resize(in_args, 1));
 
   // Create input matrix.
-  paddle_matrix mat = paddle_matrix_create(/* sample_num */ 1,
+  paddle_matrix mat = paddle_matrix_create(/* sample_num */ 10,
                                            /* size */ 784,
                                            /* useGPU */ false);
   srand(time(0));
-  paddle_real* array;
 
-  // Get First row.
-  CHECK(paddle_matrix_get_row(mat, 0, &array));
+  std::vector<paddle_real> input;
+  input.resize(784 * 10);
 
-  for (int i = 0; i < 784; ++i) {
-    array[i] = rand() / ((float)RAND_MAX);
+  for (int i = 0; i < input.size(); ++i) {
+    input[i] = rand() / ((float)RAND_MAX);
   }
+  
+  // Set value for the input matrix
+  CHECK(paddle_matrix_set_value(mat, input.data()));
 
   CHECK(paddle_arguments_set_value(in_args, 0, mat));
 
@@ -51,11 +53,17 @@ int main() {
 
   CHECK(paddle_arguments_get_value(out_args, 0, prob));
 
-  CHECK(paddle_matrix_get_row(prob, 0, &array));
+  std::std::vector<paddle_real> result;
+  int height;
+  int width;
+
+  CHECK(paddle_matrix_get_shape(prob, &height, &width);
+  result.resize(height * width);
+  CHECK(paddle_matrix_get_value(prob, result.data()));
 
   printf("Prob: ");
-  for (int i = 0; i < 10; ++i) {
-    printf("%.2f ", array[i]);
+  for (int i = 0; i < height * width; ++i) {
+    printf("%.2f ", result[i]);
   }
   printf("\n");