diff --git a/paddle/fluid/train/test_train_recognize_digits.cc b/paddle/fluid/train/test_train_recognize_digits.cc
index bd2a439f6ab5273b29010cf3599460ea8bdd68d4..45c438e8925b4e0a88e61ad509b88cd6226773a4 100644
--- a/paddle/fluid/train/test_train_recognize_digits.cc
+++ b/paddle/fluid/train/test_train_recognize_digits.cc
@@ -19,6 +19,7 @@ limitations under the License. */
 #include "gtest/gtest.h"
 
 #include "paddle/fluid/framework/executor.h"
+#include "paddle/fluid/framework/io/fs.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/framework/tensor_util.h"
diff --git a/python/paddle/fluid/tests/book/test_recognize_digits.py b/python/paddle/fluid/tests/book/test_recognize_digits.py
index 3b2c4af8ae510492052cd825077ef9dfd355b417..4fbb146752e73358a02a19fd5109e84ad00ecbae 100644
--- a/python/paddle/fluid/tests/book/test_recognize_digits.py
+++ b/python/paddle/fluid/tests/book/test_recognize_digits.py
@@ -123,8 +123,8 @@ def train(nn_type,
                     # get test acc and loss
                     acc_val = numpy.array(acc_set).mean()
                     avg_loss_val = numpy.array(avg_loss_set).mean()
-                    if float(acc_val
-                             ) > 0.2:  # Smaller value to increase CI speed
+                    if float(acc_val) > 0.2 or pass_id == (PASS_NUM - 1):
+                        # Smaller value to increase CI speed
                         if save_dirname is not None:
                             fluid.io.save_inference_model(
                                 save_dirname, ["img"], [prediction],