From 8ebc31d9358c919fdd6f50d502f4ee071a91d38e Mon Sep 17 00:00:00 2001
From: qiaolongfei <qiaolongfei@baidu.com>
Date: Wed, 4 Oct 2017 17:13:02 -0700
Subject: [PATCH] optimize the dsize

---
 paddle/operators/sgd_op.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/paddle/operators/sgd_op.h b/paddle/operators/sgd_op.h
index b501d244d7fc3..26f4012f25877 100644
--- a/paddle/operators/sgd_op.h
+++ b/paddle/operators/sgd_op.h
@@ -33,10 +33,10 @@ class SGDOpKernel : public framework::OpKernel<T> {
     auto p = framework::EigenVector<T>::Flatten(*param);
     auto g = framework::EigenVector<T>::Flatten(*grad);
     auto o = framework::EigenVector<T>::Flatten(*param_out);
-    auto lr = framework::EigenVector<T>::From(*learning_rate);
+    auto lr = framework::EigenVector<T>::Flatten(*learning_rate);
     auto place = ctx.GetEigenDevice<Place>();
 
-    Eigen::DSizes<int, 2> grad_dsize(grad->dims()[0], grad->dims()[1]);
+    Eigen::DSizes<int, 1> grad_dsize(grad->numel());
     o.device(place) = p - lr.broadcast(grad_dsize) * g;
   }
 };