May i call collate_fn within module.forward instead of dataloader

Hello, I’m wondering whether it is appropriate to collate inputs within the model, instead of the standard procedure, supplying collate_fn to dataloader.
The reason is that for a detection or segmentation model, it is better to keep inputs as List[Tensor] instead of Tensor for multi-scale images because we need the original shapes for each individual image for post-processing (e.g. clipping). The standard procedure for such case that i know of is to pad the batch images by collate_fn, and wrap the images with a structure that holds the tensor and image shapes together, like ImageList. But such custom structure as an input to module.forward is not permitted for onnx export or tracing, which prevents deployment.
Or i can input images:Tensor, image_sizes: List[Tuple[int, int]] to module.forward, but it seems redundant in deployment point of view.
A simple solution that i could think of is to keep inputs as List[Tensor], and collate within the model, not from dataloader, which allows onnx export.

What are the drawbacks of such workaround and is there any other solution for this problem?
Any help would be very much appreciated!