diff --git a/paddle/fluid/operators/gather_op.cc b/paddle/fluid/operators/gather_op.cc index 8a3450d1df97a2e99711f9ae029ca2668f38e2b0..28afeb6f541c68fe7e0719a782fd8c9147b15163 100644 --- a/paddle/fluid/operators/gather_op.cc +++ b/paddle/fluid/operators/gather_op.cc @@ -17,7 +17,7 @@ limitations under the License. */ #include #include #include "paddle/fluid/framework/ddim.h" - +#include "paddle/fluid/framework/op_version_registry.h" namespace paddle { namespace operators { @@ -152,3 +152,7 @@ REGISTER_OP_CPU_KERNEL(gather_grad, ops::GatherGradientOpKernel, ops::GatherGradientOpKernel, ops::GatherGradientOpKernel, ops::GatherGradientOpKernel); +REGISTER_OP_VERSION(gather) + .AddCheckpoint(R"ROC(upgrad gather, add attribut [axis])ROC", + paddle::framework::compatible::OpVersionDesc().NewAttr( + "axis", "Specify the axis of gather operation.", {}));