optimize the dsize

PaddlePaddle · Oct 5, 2017 · 8ebc31d · 8ebc31d
1 parent 775c602
commit 8ebc31d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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;
   }
 };