提交 0c8aeffb 编写于 作者: Y Yu Yang

Fix a PyDataProvider2 bug when use calc_batch_size

* Need PyGuard when set args for calc_batch_size
上级 3304de7a
...@@ -340,6 +340,7 @@ private: ...@@ -340,6 +340,7 @@ private:
size_t additionalBatchSize = 1; size_t additionalBatchSize = 1;
if (calcBatchSize_) { if (calcBatchSize_) {
PyGuard guard;
py::CallableHelper calcBatchSize(this->calcBatchSize_); py::CallableHelper calcBatchSize(this->calcBatchSize_);
calcBatchSize.setArgsSize(1); calcBatchSize.setArgsSize(1);
calcBatchSize.getArgs().set(0, data); calcBatchSize.getArgs().set(0, data);
...@@ -513,6 +514,7 @@ public: ...@@ -513,6 +514,7 @@ public:
} }
{ {
if (calcBatchSize_) { // custom calc batch size. if (calcBatchSize_) { // custom calc batch size.
PyGuard guard;
Py_INCREF(data.back().get()); Py_INCREF(data.back().get());
py::CallableHelper calcBatchSize(calcBatchSize_); py::CallableHelper calcBatchSize(calcBatchSize_);
calcBatchSize.setArgsSize(1); calcBatchSize.setArgsSize(1);
...@@ -575,6 +577,11 @@ public: ...@@ -575,6 +577,11 @@ public:
scanners[i]->finishFill(inArgs[i]); scanners[i]->finishFill(inArgs[i]);
} }
{
PyGuard g;
cache_->drop(&data);
}
DBG << "Reading CPU Batch Done."; DBG << "Reading CPU Batch Done.";
if (useGpu_) { if (useGpu_) {
...@@ -592,10 +599,6 @@ public: ...@@ -592,10 +599,6 @@ public:
*batch = cpuBatch; *batch = cpuBatch;
} }
{
PyGuard g;
cache_->drop(&data);
}
return bsize; return bsize;
} }
}; };
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册