From 0c8aeffb9acd1d727d4f1c4bb18d4c6ce3346891 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Fri, 9 Sep 2016 17:01:20 +0800
Subject: [PATCH] Fix a PyDataProvider2 bug when use calc_batch_size

* Need PyGuard when set args for calc_batch_size
---
 paddle/gserver/dataproviders/PyDataProvider2.cpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/paddle/gserver/dataproviders/PyDataProvider2.cpp b/paddle/gserver/dataproviders/PyDataProvider2.cpp
index f7886c4e014..8e51752dc29 100644
--- a/paddle/gserver/dataproviders/PyDataProvider2.cpp
+++ b/paddle/gserver/dataproviders/PyDataProvider2.cpp
@@ -340,6 +340,7 @@ private:
 
       size_t additionalBatchSize = 1;
       if (calcBatchSize_) {
+        PyGuard guard;
         py::CallableHelper calcBatchSize(this->calcBatchSize_);
         calcBatchSize.setArgsSize(1);
         calcBatchSize.getArgs().set(0, data);
@@ -513,6 +514,7 @@ public:
       }
       {
         if (calcBatchSize_) {  // custom calc batch size.
+          PyGuard guard;
           Py_INCREF(data.back().get());
           py::CallableHelper calcBatchSize(calcBatchSize_);
           calcBatchSize.setArgsSize(1);
@@ -575,6 +577,11 @@ public:
       scanners[i]->finishFill(inArgs[i]);
     }
 
+    {
+      PyGuard g;
+      cache_->drop(&data);
+    }
+
     DBG << "Reading CPU Batch Done.";
 
     if (useGpu_) {
@@ -592,10 +599,6 @@ public:
       *batch = cpuBatch;
     }
 
-    {
-      PyGuard g;
-      cache_->drop(&data);
-    }
     return bsize;
   }
 };
-- 
GitLab