I just tested this and it works. I replace Flatten in the create_head function with your definition and exported to ONNX, then imported the model in onnx.js and it works the same as if I used nn.Flatten directly (as expected).
I just tested this and it works. I replace Flatten in the create_head function with your definition and exported to ONNX, then imported the model in onnx.js and it works the same as if I used nn.Flatten directly (as expected).