single_op

This commit is contained in:
yefeng 2021-04-21 09:56:34 +08:00
parent c37f8f9f0d
commit 39e3f11b84
2 changed files with 3 additions and 2 deletions

View File

@ -123,8 +123,6 @@ int TileCPUKernel::SimpleTileImpl(int task_id) {
}
int TileCPUKernel::RunSimpleTile() {
auto data_type = in_tensors_.at(0)->data_type();
tile_parameter_->data_size_ = lite::DataTypeSize(data_type);
auto ret = ParallelLaunch(static_cast<const lite::InnerContext *>(this->context_)->thread_pool_, SimpleTile, this,
context_->thread_num_);
if (ret != RET_OK) {
@ -135,6 +133,8 @@ int TileCPUKernel::RunSimpleTile() {
}
int TileCPUKernel::Run() {
auto data_type = in_tensors_.at(0)->data_type();
tile_parameter_->data_size_ = lite::DataTypeSize(data_type);
input_addr_ = reinterpret_cast<uint8_t *>(in_tensors_.at(0)->data_c());
output_addr_ = reinterpret_cast<uint8_t *>(out_tensors_.at(0)->data_c());
MS_ASSERT(input_addr_ != nullptr);

View File

@ -37,6 +37,7 @@ ops::PrimitiveC *TFArgMinParser::Parse(const tensorflow::NodeDef &tf_op,
auto &axis_tensor = attr_value.tensor();
prim->set_axis(axis_tensor.int_val(0));
prim->set_out_max_value(false);
prim->set_top_k(1);
*output_size = 1;
if (AddOpInput(tf_op, 0, inputs) != RET_OK) {