Fix locking of variables in SparseProximalGradientDescent,

AdagradDA, SparseAdagradDA. PiperOrigin-RevId: 168252530
tensorflow · Sep 11, 2017 · a6223c0 · a6223c0
1 parent abde008
commit a6223c0
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc
@@ -875,7 +875,7 @@ class SparseApplyProximalGradientDescentOp : public OpKernel {
 
   void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
     auto locks =
-        MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+        MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0});
     Tensor var;
     OP_REQUIRES_OK(
         ctx, GetInputTensorFromVariable(ctx, 0, use_exclusive_lock_, &var));
@@ -1524,8 +1524,8 @@ class ApplyAdagradDAOp : public OpKernel {
   }
 
   void Compute(OpKernelContext* ctx) override {
-    auto locks =
-        MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+    auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
+                                                      {0, 1, 2});
     Tensor var;
     OP_REQUIRES_OK(
         ctx, GetInputTensorFromVariable(ctx, 0, use_exclusive_lock_, &var));
@@ -1626,8 +1626,8 @@ class SparseApplyAdagradDAOp : public OpKernel {
   }
 
   void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS {
-    auto locks =
-        MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_, {0, 1});
+    auto locks = MaybeLockVariableInputMutexesInOrder(ctx, use_exclusive_lock_,
+                                                      {0, 1, 2});
     Tensor var;
     OP_REQUIRES_OK(
         ctx, GetInputTensorFromVariable(ctx, 0, use_exclusive_lock_, &var));