fix: bug in elementwise base for static inputs by zewenli98 · Pull Request #2819 · pytorch/TensorRT (original) (raw)

Could this be the error you mentioned occurring in TRT-10?

@chohk88 Yes that's it! The issue here is that the original broadcast() function in fx/conterver_utils.py only broadcasts two ITensors to the same number of dimensions (ranks) by prepending 1s, which results in the error in the following ctx.net.add_elementwise(). For fixing this, I wrote another similar broadcast function, called broadcast_to_same_shape that can broadcast two ITensors to the exactly same shape.

For example, we have original ITensors: lhs_val.shape: (2, 3) rhs_val.shape: (2, 2, 1, 3)

• If calling fx/converter_utils.broadcast, lhs_val.shape: (1, 1, 2, 3) lhs_val.shape: (2, 2, 1, 3).
• If calling broadcast_to_same_shape, lhs_val.shape: (2, 2, 2, 3) lhs_val.shape: (2, 2, 2, 3).

I think this fix should also work for #2726