Fix generator not closed when iteration end bug

7 years ago · e7d87c09fc
parent 515a756a3c
commit e7d87c09fc
1 changed files with 2 additions and 0 deletions
--- a/benchmark/fluid/fluid_benchmark.py
+++ b/benchmark/fluid/fluid_benchmark.py
@ -151,6 +151,7 @@ def train(avg_loss, infer_prog, optimizer, train_reader, test_reader, batch_acc,
                if data == None:
                    break
            if iters == args.iterations:
+                reader_generator.close()
                break
            if iters == args.skip_batch_num:
                start_time = time.time()
@ -252,6 +253,7 @@ def train_parallel(avg_loss, infer_prog, optimizer, train_reader, test_reader,
                if data == None:
                    break
            if iters == args.iterations:
+                reader_generator.close()
                break
            if args.profile and pass_id == 0 and batch_id == 5:
                profiler.start_profiler("All")