Fix RandomShuffleQueue initialization to reserve all queue components.

Change: 133622066

1 files changed, 2 insertions, 3 deletions

diff --git a/tensorflow/core/kernels/random_shuffle_queue_op.cc b/tensorflow/core/kernels/random_shuffle_queue_op.cc
index 6514e1fd45..6b3b792814 100644
--- a/tensorflow/core/kernels/random_shuffle_queue_op.cc
+++ b/tensorflow/core/kernels/random_shuffle_queue_op.cc
@@ -102,12 +102,11 @@ RandomShuffleQueue::RandomShuffleQueue(
 }
 
 Status RandomShuffleQueue::Initialize() {
-  Status s = TypedQueue::Initialize();
-  if (!s.ok()) return s;
+  TF_RETURN_IF_ERROR(TypedQueue::Initialize());
 
   mutex_lock lock(mu_);
   for (int i = 0; i < num_components(); ++i) {
-    queues_.back().reserve(min_after_dequeue_);
+    queues_[i].reserve(min_after_dequeue_);
   }
   return Status::OK();
 }