How does mmdeploy convert FCOS3D's postprocessing pipeline to ONNX?

[tensorturtle]

I built a PyTorch model from the components of FCOS3D as follows:

I found the forward pass top level at: https://github.com/open-mmlab/mmdetection3d/blob/47285b3f1e9dba358e98fcd12e523cfd0769c876/mmdet3d/models/detectors/single_stage_mono3d.py#L110-L113

import torch 

config_file = '/home/ubuntu/Repos/mmdetection3d/configs/fcos3d/fcos3d_r101_caffe_fpn_gn-head_dcn_2x8_1x_nus-mono3d.py'
checkpoint_file = '/home/ubuntu/Repos/mmdetection3d/work_dirs/fcos3d_r101_caffe_fpn_gn-head_dcn_2x8_1x_nus-mono3d_finetune_20210717_095645-8d806dc2.pth'
model = init_model(config_file, checkpoint_file, device='cuda:0')

# reconstruct model by decomposing its parts
backbone = model.backbone
neck = model.neck
bbox_head = model.bbox_head

class PureModel(torch.nn.Module):
    def __init__(self, backbone, neck, bbox_head):
        super().__init__()
        self.backbone = backbone
        self.neck = neck
        self.bbox_head = bbox_head

    def forward(self, x):
        x = self.backbone(x)
        x = self.neck(x)
        x = self.bbox_head(x)
        # The following line contains operations that are not supported by ONNX:
        x = self.bbox_head.get_bboxes(*x, img_metas=data['img_metas'][0])
        return x

pure_model = PureModel(backbone, neck, bbox_head)

When I run torch.onnx.export on the above pure_model WITHOUT the last self.bboxes_head.get_bboxes (postprocessing step), it exports successfully and can further be converted into TensorRT with seemingly correct results.

However, when I include the self.bboxes_head.get_bboxes, ONNX export is not successful. There are many TracerWarnings from attempting to convert tensor to Python boolean, and also:

RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: CameraInstance3DBoxes

So currently, I have to split up the inference into two steps (which also means transferring data back and forth to the GPU twice)

My question is, how is mmdeploy able to convert FCOS3D into a single-step ONNX model, when it seems to contain several incompatible tensor operations in the final step?

Thank you