Add support for exporting with aux task ouputs

This adds the functionality to retrieve the output of the "track origin" aux task from the exported onnx model.

Some code for exporting the output of the other aux task, i.e. track vertexing, was already in main branch. This MR modifies that code as well to make the final code more concise (and also fixes an edge-case where squeeze() converted [[0]] to scalar 0).