Fix rank_attention op_version, test=op_version (#30006)

* fix rank_attention, test=op_version
revert-31562-mean
ShenLiang 4 years ago committed by GitHub
parent 2bc5121da8
commit 893d37e5c6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -13,6 +13,7 @@ limitations under the License. */
#include <memory>
#include <string>
#include <vector>
#include "paddle/fluid/framework/op_version_registry.h"
namespace paddle {
namespace operators {
@ -176,3 +177,18 @@ REGISTER_OP_CPU_KERNEL(
rank_attention,
ops::RankAttentionKernel<paddle::platform::CPUDeviceContext, float>,
ops::RankAttentionKernel<paddle::platform::CPUDeviceContext, double>);
REGISTER_OP_VERSION(rank_attention)
.AddCheckpoint(
R"ROC(
Upgrade rank_attention, add 1 outputs [InputHelp] and 1 attribute
[MaxSize].
)ROC",
paddle::framework::compatible::OpVersionDesc()
.NewOutput("InputHelp",
"Output tensor of rank_attention_Op operator "
"in order to assist calculation in the reverse process.")
.NewAttr(
"MaxSize",
"Forward calculation to set the pre-applied video memory size",
0));

Loading…
Cancel
Save