AttributeError: "DecoderOnlyLanguageModel" object has no attribute "num_heads" when I use "alphageometry" mode?

How to fix this problem when I use alphageometry mode?
Traceback (most recent call last):
File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
exec(code, run_globals)
File "/content/alphageometry/alphageometry.py", line 645, in
app.run(main)
File "/usr/local/lib/python3.10/dist-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/usr/local/lib/python3.10/dist-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/content/alphageometry/alphageometry.py", line 631, in main
model = get_lm(_CKPT_PATH.value, _VOCAB_PATH.value)
File "/content/alphageometry/alphageometry.py", line 203, in get_lm
return lm.LanguageModelInference(vocab_path, ckpt_init, mode='beam_search')
File "/content/alphageometry/lm_inference.py", line 66, in __init__
self.n = imodel.num_heads
File "/usr/local/lib/python3.10/dist-packages/flax/linen/module.py", line 718, in __getattr__
raise AttributeError(
AttributeError: "DecoderOnlyLanguageModel" object has no attribute "num_heads"