TF2.2 PB file conversion to ONNX "<function rewrite_single_direction_lstm" failure

Description

I have trained a TF2.2 model, which i am trying to convert to onnx. it gives me “<function rewrite_single_direction_lstm” failure. please find saved_model and TF2.2.0 PB file.

Environment

TensorRT Version : 7.0.0.11
GPU Type : Tesla K80
Nvidia Driver Version : 450.36.06
CUDA Version : CUDA Version: 11.0
CUDNN Version : 8.0
Operating System + Version : Ubuntu x86_64
Python Version (if applicable) : 3.7.7
TensorFlow Version (if applicable) : 2.2.0
PyTorch Version (if applicable) :
Baremetal or Container (if container which image + tag) : barematel

Relevant Files

Saved_model:

Froze graph:

Steps To Reproduce

(tensor2) ubuntu@usva-asd-ml-ear1:~/notebook/cnn-facial-landmark$ python -m tf2onnx.convert --graphdef frozen_graph_t2.pb --output frozen_graph.onnx --inputs image_tensor:0 --outputs layer6/final_dense:0 --fold_const --opset 11

2020-06-09 15:13:01,934 - INFO - Froze 0 variables.
INFO:tensorflow:Converted 0 variables to const ops.
2020-06-09 15:13:01,948 - INFO - Converted 0 variables to const ops.
2020-06-09 15:13:02,322 - INFO - Using tensorflow=2.2.0, onnx=1.7.0, tf2onnx=1.6.1/aafc83
2020-06-09 15:13:02,322 - INFO - Using opset <onnx, 11>
2020-06-09 15:13:03,327 - ERROR - rewriter <function rewrite_single_direction_lstm at 0x7f75c1a27cb0>: exception (‘switch_true map/while/map/while/loop_counter_switch/_15 has unexpected count of consumers:’, [‘map/while/loop_body_control/_21’, ‘Func/map/while/body/_1/input/_44’])

Hi @yaduvir.singh,

I would recommend filing an issue under tf2onnx here: https://github.com/onnx/tensorflow-onnx/issues.