diff options
| -rw-r--r-- | tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h | 7 |
1 files changed, 3 insertions, 4 deletions
diff --git a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h index a0606427a5..6ed177e001 100644 --- a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h +++ b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h @@ -399,7 +399,7 @@ ASBSQueue<TaskType>::~ASBSQueue() { template <typename TaskType> Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) { - bool added_new_batch = false; + ASBSBatch<TaskType>* new_batch = nullptr; size_t size = (*task)->size(); if (size > options_.max_batch_size) { return errors::InvalidArgument("Task size ", size, @@ -418,15 +418,14 @@ Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) { current_batch_ = nullptr; } if (!current_batch_) { - added_new_batch = true; num_enqueued_batches_++; - current_batch_ = + current_batch_ = new_batch = new ASBSBatch<TaskType>(this, scheduler_->GetEnv()->NowMicros()); } current_batch_->AddTask(std::move(*task)); num_enqueued_tasks_++; } - if (added_new_batch) scheduler_->AddBatch(current_batch_); + if (new_batch != nullptr) scheduler_->AddBatch(new_batch); return Status::OK(); } |