TensorRT problem in PX 2 (example attached)

I am facing a TensorRT-related bug, which appears only on PX 2 (OS 5.0.10.3), and not on Ubuntu (TensorRT 4.0.1.6).
This is a simple case: when a UFF plugin layer has two or more input.

A network is this simple:

with tf.Graph().as_default():
    
    image_ph = tf.placeholder(tf.float32, [1, 5, 5, 3])
    net = slim.conv2d(image_ph, 3, [3, 3])
    branch0 = slim.conv2d(net, 3, [3, 3])
    branch1 = slim.conv2d(net, 3, [3, 3])
            
    def merge(b1, b2): return b1 + 1    
    output = tf.py_func(merge, [branch0, branch1], tf.float32, name="output")  # custom layer needed!

I attached its frozen graph (why.pb), UFF (why.pb.uff) and Pbtxt (why.pb.uff.pbtxt).
FYI, to replace custom layer, I used “uff-surgery-why.py” (attached).

convert-to-uff tensorflow --input-file why.pb -O output -p uff-surgery-why.py -t

** why.pb, why.pb.uff, why.pb.pbtxt, uff-surgery-why.py are in “why.zip”

I also attached a simple C++ project running inference. “px2-bug-minimal.zip”.
You build the project by unzipping and doing “mkdir build && cd build && cmake … && make -j8 && ./px2-bug-minimal”. (px2-bug-minimal.cpp should have a path to why.pb.uff)

When running on Unbuntu, it works fine.
When running on PX 2,

INFO: UFFParser: parsing Conv/Relu
INFO: UFFParser: parsing Conv_1/weights
INFO: UFFParser: parsing Conv_1/Conv2D
INFO: UFFParser: parsing Conv_1/biases
INFO: UFFParser: parsing Conv_1/BiasAdd
INFO: UFFParser: parsing Conv_1/Relu
INFO: UFFParser: parsing Conv_2/weights
INFO: UFFParser: parsing Conv_2/Conv2D
INFO: UFFParser: parsing Conv_2/biases
INFO: UFFParser: parsing Conv_2/BiasAdd
INFO: UFFParser: parsing Conv_2/Relu
INFO: UFFParser: parsing output
terminate called after throwing an instance of 'std::out_of_range'                                   
  what():  _Map_base::at
Aborted (core dumped)

Is this because TensorRT version in the newest PX 2 OS is still older than Ubuntu TensorRT 4.0.1.6? Or am I missing something? Any hints or ideas? Please let me know.

why.zip (4.71 KB)
px2-bug-minimal.zip (17.1 KB)

We will support you via below topic:
https://devtalk.nvidia.com/default/topic/1044208/driveworks/px2-tensorrt-problem/