fix the GetKernelTypeForVar of input for fluid.gather (#28534)

4 years ago · c52fe48f6f
parent 621b31c526
commit c52fe48f6f
1 changed files with 5 additions and 1 deletions
--- a/paddle/fluid/operators/gather_op.cc
+++ b/paddle/fluid/operators/gather_op.cc
@ -69,7 +69,11 @@ class GatherOp : public framework::OperatorWithKernel {
  framework::OpKernelType GetKernelTypeForVar(
      const std::string& var_name, const framework::Tensor& tensor,
      const framework::OpKernelType& expected_kernel_type) const override {
-    return expected_kernel_type;
+    if (var_name == "Axis") {
+      return expected_kernel_type;
+    }
+    return framework::OpKernelType(expected_kernel_type.data_type_,
+                                   tensor.place(), tensor.layout());
  }
 };