提交 d11ba0d2 编写于 作者: K Kavya Srinet

Fixed assign_value_op.cc

上级 e3351705
...@@ -36,7 +36,7 @@ class AssignValueOp : public framework::OperatorWithKernel { ...@@ -36,7 +36,7 @@ class AssignValueOp : public framework::OperatorWithKernel {
framework::OpKernelType GetExpectedKernelType( framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext &ctx) const override { const framework::ExecutionContext &ctx) const override {
return framework::OpKernelType( return framework::OpKernelType(
framework::proto::DataType(ctx.Attr<int>("dtype")), ctx.GetPlace()); framework::proto::VarType::Type(ctx.Attr<int>("dtype")), ctx.GetPlace());
} }
}; };
...@@ -49,8 +49,8 @@ class AssignValueOpMaker : public framework::OpProtoAndCheckerMaker { ...@@ -49,8 +49,8 @@ class AssignValueOpMaker : public framework::OpProtoAndCheckerMaker {
"(vector<int>) " "(vector<int>) "
"Shape of values."); "Shape of values.");
AddAttr<int>("dtype", "data type of values") AddAttr<int>("dtype", "data type of values")
.InEnum({framework::proto::DataType::INT32, .InEnum({framework::proto::VarType::Type::INT32,
framework::proto::DataType::FP32}); framework::proto::VarType::Type::FP32});
AddAttr<std::vector<float>>("fp32_values", "store the float values") AddAttr<std::vector<float>>("fp32_values", "store the float values")
.SetDefault({}); .SetDefault({});
AddAttr<std::vector<int>>("int32_values", "store the int values") AddAttr<std::vector<int>>("int32_values", "store the int values")
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册