forked from mindspore-Ecosystem/mindspore
single_op
This commit is contained in:
parent
c37f8f9f0d
commit
39e3f11b84
|
@ -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);
|
||||
|
|
|
@ -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) {
|
||||
|
|
Loading…
Reference in New Issue