Address @elliottslaughter's idea in https://reviewable.io/reviews/flexflow/flexflow-train/1629#-OmOVcMl4gGR8Ink49vR