ONNX export for train time model

[YevheniiStarship]

Hello! Didn't find similar topic, however, I might miss something. While I am exporting yolov5 model I want get rid of inference part and return only output of last three convolution layers. I will attach image of what I mean. Is it possible to achieve it with existing export_onnx function? Would be glad of any help! Thank you!