diff --git a/paddle/framework/operator.cc b/paddle/framework/operator.cc index 946bde5734ba37ebaa8d2e142b3c199fda2a397d..1a7e33222788e5751b4049bac89f55aa9dd4925c 100644 --- a/paddle/framework/operator.cc +++ b/paddle/framework/operator.cc @@ -18,14 +18,14 @@ namespace paddle { namespace framework { template <> -Eigen::DefaultDevice* OpKernel::KernelContext::get_eigen_device< +Eigen::DefaultDevice* OpKernel::KernelContext::GetEigenDevice< platform::CPUPlace, Eigen::DefaultDevice>() const { return device_context_.get_eigen_device(); } #ifndef PADDLE_ONLY_CPU template <> -Eigen::GpuDevice* OpKernel::KernelContext::get_eigen_device< +Eigen::GpuDevice* OpKernel::KernelContext::GetEigenDevice< platform::GPUPlace, Eigen::GpuDevice>() const { return device_context_.get_eigen_device(); } diff --git a/paddle/framework/operator.h b/paddle/framework/operator.h index e6cae9c32b372e27b62d521afbb07f0396f8d6b9..b8c5098e49a7118994fe8132027953f6180ac624 100644 --- a/paddle/framework/operator.h +++ b/paddle/framework/operator.h @@ -109,7 +109,9 @@ class OpKernel { template ::EigenDeviceType> - DeviceType* get_eigen_device() const; + DeviceType* GetEigenDevice() const; + + platform::Place GetPlace() const { return device_context_.GetPlace(); } const OperatorBase& op_; const ScopePtr& scope_; diff --git a/paddle/operators/add_op.h b/paddle/operators/add_op.h index e8c718669a9959252ce473e989a0ad27ebc487cc..e9a793d23b3cb4dab42eb2985f792d85f181420d 100644 --- a/paddle/operators/add_op.h +++ b/paddle/operators/add_op.h @@ -27,9 +27,9 @@ public: auto input1 = context.Input(1)->Get(); auto* output = context.Output(0)->GetMutable(); - output->mutable_data(Place()); + output->mutable_data(context.GetPlace()); - output->flat().device(*(context.get_eigen_device())) = + output->flat().device(*(context.GetEigenDevice())) = input0.flat() + input1.flat(); } };