From 698b8b73ad62f3af069b8b486efb070b2afe5c39 Mon Sep 17 00:00:00 2001
From: Zhang Ting <709968123@qq.com>
Date: Fri, 22 Nov 2019 11:15:15 +0800
Subject: [PATCH] optimize lod_reset op to avoid data transform

---
 paddle/fluid/operators/lod_reset_op.cc | 9 +++++++++
 paddle/fluid/operators/lod_reset_op.h  | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/lod_reset_op.cc b/paddle/fluid/operators/lod_reset_op.cc
index baa709b9e4..94483bc7f4 100644
--- a/paddle/fluid/operators/lod_reset_op.cc
+++ b/paddle/fluid/operators/lod_reset_op.cc
@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/lod_reset_op.h"
 #include <memory>
+#include <string>
 
 namespace paddle {
 namespace operators {
@@ -50,6 +51,14 @@ class LoDResetOp : public framework::OperatorWithKernel {
         OperatorWithKernel::IndicateVarDataType(ctx, "X"),
         ctx.device_context());
   }
+
+  framework::OpKernelType GetKernelTypeForVar(
+      const std::string &var_name, const framework::Tensor &tensor,
+      const framework::OpKernelType &expected_kernel_type) const override {
+    return framework::OpKernelType(expected_kernel_type.data_type_,
+                                   expected_kernel_type.place_,
+                                   tensor.layout());
+  }
 };
 
 class LoDResetOpVarTypeInference : public framework::VarTypeInference {
diff --git a/paddle/fluid/operators/lod_reset_op.h b/paddle/fluid/operators/lod_reset_op.h
index d827f2a2eb..7677fa2251 100644
--- a/paddle/fluid/operators/lod_reset_op.h
+++ b/paddle/fluid/operators/lod_reset_op.h
@@ -45,7 +45,7 @@ class LoDResetKernel : public framework::OpKernel<T> {
         return;  // early return, since lod already set
       } else {
         auto* lod = lod_t->data<int>();
-        if (platform::is_gpu_place(ctx.GetPlace())) {
+        if (platform::is_gpu_place(lod_t->place())) {
           framework::Tensor lod_cpu;
           framework::TensorCopySync(*lod_t, platform::CPUPlace(), &lod_cpu);
           lod = lod_cpu.data<int>();
-- 
GitLab