TensorRT转换onnx的Transpose算子遇到的奇怪问题

近来把一个模型导出为onnx并用onnx simplifier化简后转换为TensorRT engine遇到非常奇怪的问题，在我们的网络中有多个检测头时，转换出来的engine的推理效果是正常的，当网络中只有一个检测头时，转换出来的engine的推理效果奇差，类别得分经sigmoid()计算出来的很多很大的值，导致输出上百万个3D box，这明显不正常。

开始以为3D box的C++后处理解码部分有问题，调查折腾了不少时间发现根本没什么大问题，然后想起模型只有一个检测头输出不正常，有多个头时检测头时反而正常，于是去仔细翻看模型网络的实现代码，也没发现什么不对的，很郁闷，想到把用于TensorRT生成engine的onnx可视化出来看有何差异，刚开始看也没发现什么特别的，感觉展示出来的结构都很正确没毛病啊：

那为何多个检测头就正常一个检测头就不行呢？再仔细看代码，发现heatmap这个分支的输出tensor在作permute（onnx里对应的是Transpose）后有这么一句有关：

      result['hm_preds'] =  torch.cat(hm_preds_list, dim=-1)

hm_preds_list是个list，当有多个head时，会有多个hm_preds加入到这个[ ]里，只有一个检测头时就只有一个hm_preds加入这个[ ]里，很显然，当hm_preds_list里只有一个hm_preds时，torch.cat()不会作实质操作，只是依据hm_preds_list里的数据直接输出tensor而已，那我强制只有一个检测头时