|
|
|
@ -35,14 +35,15 @@ class ShapeOp : public framework::OperatorWithKernel {
|
|
|
|
|
class ShapeOpMaker : public framework::OpProtoAndCheckerMaker {
|
|
|
|
|
public:
|
|
|
|
|
void Make() override {
|
|
|
|
|
AddInput("Input", "(Tensor), The input tensor.");
|
|
|
|
|
AddOutput("Out",
|
|
|
|
|
"(Tensor), The shape of input tensor, the data type of the shape"
|
|
|
|
|
" is int32_t, will be on the same device with the input Tensor.");
|
|
|
|
|
AddInput("Input", "(LoDTensor), The input tensor.");
|
|
|
|
|
AddOutput(
|
|
|
|
|
"Out",
|
|
|
|
|
"(LoDTensor), The shape of input tensor, the data type of the shape"
|
|
|
|
|
" is int32_t, will be on the same device with the input Tensor.");
|
|
|
|
|
AddComment(R"DOC(
|
|
|
|
|
Shape Operator
|
|
|
|
|
Shape Operator.
|
|
|
|
|
|
|
|
|
|
Get the shape of input tensor. Only support CPU input Tensor now.
|
|
|
|
|
Return the shape of the input.
|
|
|
|
|
)DOC");
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|