config.json

{
  "model_type": "esa",
  "auto_map": {
    "AutoConfig": "configuration_esa.ESAConfig",
    "AutoModel": "modeling_esa.ESAModel"
  },
  "node_feat_dim": 79,
  "edge_feat_dim": 13,
  "apply_attention_on": "edge",
  "hidden_dims": [256, 256, 256, 256, 256],
  "num_heads": [4, 4, 4, 4, 4],
  "layer_types": ["M", "S", "M", "S", "P"],
  "dim_output": 256,
  "num_pma_seeds": 32,
  "use_mlps": true,
  "mlp_type": "gated_mlp",
  "mlp_hidden_size": 64,
  "num_mlp_layers": 2,
  "use_mlp_ln": true,
  "mlp_dropout": 0.0,
  "xformers_or_torch_attn": "torch",
  "pre_or_post": "post",
  "norm_type": "LN",
  "sab_dropout": 0.0,
  "mab_dropout": 0.0,
  "pma_dropout": 0.0,
  "residual_dropout": 0.0,
  "pma_residual_dropout": 0.0,
  "use_bfloat16": true,
  "set_max_items": 0
}