{
  "seed": 42,
  "wandb": true,
  "truncation_seq_length": 512,
  "batch_size": 16,
  "datamodule": {
    "batch_size": 64,
    "num_workers": 32,
    "_target_": "foldmae.data.TEDLightningDataset",
    "root": "./datasets/afdb_FS_plddt80",
    "dataset_name": "ted"
  },
  "trainer": {
    "_target_": "pytorch_lightning.Trainer",
    "accelerator": "auto",
    "max_steps": 18300,
    "strategy": "ddp_find_unused_parameters_true",
    "devices": "auto",
    "default_root_dir": "${logs.path}",
    "num_sanity_val_steps": 0,
    "accumulate_grad_batches": 16
  },
  "model": {
    "name": "esm2_t33_650M_UR50D",
    "num_classes": 965,
    "avg_pool": false
  },
  "train": {
    "optimizer": {
      "_target_": "torch.optim.AdamW",
      "lr": 0.0016,
      "weight_decay": 0.1,
      "betas": [
        0.9,
        0.95
      ]
    },
    "lr_scheduler": {
      "_target_": "foldmae.lr_schedulers.get_cosine_schedule_with_warmup",
      "warmup_steps": 1830,
      "max_steps": "${trainer.max_steps}"
    },
    "loss": {
      "_target_": "torch.nn.CrossEntropyLoss"
    },
    "llrd": 0.8,
    "ckpt_path": "logs/esm2/finetune/ted/42/runs/2026-01-20_07-59-12/model.ckpt"
  },
  "logs": {
    "prefix": "logs/esm2/finetune/${datamodule.dataset_name}/${seed}",
    "path": "${logs.prefix}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S}"
  },
  "_model_class": "esm2"
}