optimize lod_reset op to avoid data transform

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 {

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>();