[TF FE] Fix BatchToSpace translator (#12815) (#12834)

According to the specification we must have the same type for block_shape and crops inputs

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
This commit is contained in:
Roman Kazantsev 2022-08-31 18:17:29 +03:00 committed by GitHub
parent d9d4b6d89b
commit c6d1d2aed1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -69,12 +69,15 @@ OutputVector translate_batch_nd_and_space_nd_op(const NodeContext& node) {
auto crops_begin = make_shared<Squeeze>(crops_split->outputs()[0], axes);
auto crops_end = make_shared<Squeeze>(crops_split->outputs()[1], axes);
// types of block_shape and crops inputs must be the same according to the specification
auto converted_padded_block_shape = make_shared<Convert>(padded_block_shape, crops_begin->get_element_type());
if (node.get_op_type() == "BatchToSpaceND") {
auto res = make_shared<BatchToSpace>(input, padded_block_shape, crops_begin, crops_end);
auto res = make_shared<BatchToSpace>(input, converted_padded_block_shape, crops_begin, crops_end);
set_node_name(node.get_name(), res);
return res->outputs();
} else if (node.get_op_type() == "SpaceToBatchND") {
auto res = make_shared<SpaceToBatch>(input, padded_block_shape, crops_begin, crops_end);
auto res = make_shared<SpaceToBatch>(input, converted_padded_block_shape, crops_begin, crops_end);
set_node_name(node.get_name(), res);
return res->outputs();
}