{
  "device": "mps",
  "samples": 4,
  "parameters": 817346,
  "initial": {
    "stage": "pretrain",
    "step": 1,
    "loss": 80.38571166992188,
    "text_loss": 80.38571166992188,
    "flow_loss": null
  },
  "final": {
    "stage": "posttrain",
    "step": 500,
    "loss": 0.1929057091474533,
    "text_loss": 2.784368575703411e-07,
    "flow_loss": 0.09645271301269531
  }
}