config.py

import ml_collections as mlc

consts = mlc.ConfigDict(
    {
        "batch_size": 2,
        "n_res": 11,
        "n_seq": 13,
        "n_templ": 3,
        "n_extra": 17,
        "eps": 5e-4,
        # For compatibility with DeepMind's pretrained weights, it's easiest for
        # everyone if these take their real values.
        "c_m": 256,
        "c_z": 128,
        "c_s": 384,
        "c_t": 64,
        "c_e": 64,
    }
)