Train GNN on Upgrade MC

Created by: asogaard

Is blocked by #46