pretrain_config.json

{
  "dataset_name": "METR-LA",
  "pretrain_config": {
    "stage1_epochs": 1,
    "stage1_mask_ratio": 0.15,
    "stage2_epochs": 1,
    "stage2_mask_ratio": 0.1,
    "learning_rate": 0.001,
    "pretrain_batch_size": 100,
    "imputation_iterations": 3,
    "save_to": "STGFORMER_TEST_E2E"
  },
  "model_dim": 128,
  "num_nodes": 207
}