diff --git a/src/operator/TopK.cpp b/src/operator/TopK.cpp index 52bb3755431b32cc3d30a85507e0e3fa22e0250c..660865100a817b32c42328b34032541479b3aefc 100644 --- a/src/operator/TopK.cpp +++ b/src/operator/TopK.cpp @@ -38,7 +38,10 @@ TopK_Op::TopK_Op( attr<TopKAttr::Axis>(axis), attr<TopKAttr::Largest>(largest), attr<TopKAttr::Sorted>(sorted), - attr<TopKAttr::K>(k))) {} + attr<TopKAttr::K>(k))) +{ + mOutputs[1]->setDataType(DataType::Int64); +} TopK_Op::TopK_Op(const TopK_Op& op) : OperatorTensor(op), @@ -79,7 +82,6 @@ bool Aidge::TopK_Op::forwardDims(bool allowDataDependency) { outDims[kAxis] = k(); mOutputs[0]->resize(outDims); mOutputs[1]->resize(outDims); - mOutputs[1]->setDataType(DataType::Int64); return true; }