{
  "device": "mps",
  "samples": 4,
  "parameters": 816482,
  "initial": {
    "stage": "pretrain",
    "step": 1,
    "loss": 78.87449645996094,
    "text_loss": 78.87449645996094,
    "flow_loss": null
  },
  "final": {
    "stage": "posttrain",
    "step": 500,
    "loss": 0.12353827804327011,
    "text_loss": 3.1691649837739533e-06,
    "flow_loss": 0.06176755577325821
  }
}