IntenseWebs/openvino
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix Unpack translator in TF FE (#10494)

Browse Source 
* Fix Unpack translator in TF FE

* Apply review feedback
This commit is contained in:
Maxim Vafin
2022-02-19 02:52:48 +03:00
committed by GitHub
parent 2e164b4ddc
commit 71fdcdf899
1 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/frontends/tensorflow/src/op/unpack.cpp
View File

@@ -19,9 +19,17 @@ OutputVector translate_unpack_op(const NodeContext& node) {
auto num = node.get_attribute<int64_t>("num");
auto axis_const = make_shared<Constant>(element::i64, Shape{}, axis);
auto res = make_shared<Split>(input, axis_const, num);
set_node_name(node.get_name(), res);
return res->outputs();
auto split = make_shared<Split>(input, axis_const, num);
OutputVector res;
int idx = 0;
for (auto out : split->outputs()) {
auto squeezed_res = make_shared<Squeeze>(out, axis_const);
squeezed_res->set_friendly_name(node.get_name() + "/squeeze_" + to_string(idx));
set_out_name(node.get_name() + ":" + std::to_string(idx), squeezed_res->output(0));
++idx;
res.push_back(squeezed_res);
}
return res;
}
} // namespace op
} // namespace tensorflow