Already done in some other manifestation of this template (heptrkx-gnn-tracking). I just need to port it here.