cuEq-accelerated MACE gives error for num_interactions=1

[lucasdekam]

Hi, when running MACE (installed from source, not pip) with num_interactions=1 and enable_cueq=True I get the following error.

2024-12-13 13:07:11.655 INFO: Converting model to CUEQ for accelerated training
Traceback (most recent call last):
  File "/home/ldkam/miniconda3/envs/cueq/bin/mace_run_train", line 8, in <module>
    sys.exit(main())
  File "/home/ldkam/miniconda3/envs/cueq/lib/python3.10/site-packages/mace/cli/run_train.py", line 66, in main
    run(args)
  File "/home/ldkam/miniconda3/envs/cueq/lib/python3.10/site-packages/mace/cli/run_train.py", line 581, in run
    model = run_e3nn_to_cueq(deepcopy(model), device=device)
  File "/home/ldkam/miniconda3/envs/cueq/lib/python3.10/site-packages/mace/cli/convert_e3nn_cueq.py", line 132, in run
    config = extract_config_mace_model(source_model)
  File "/home/ldkam/miniconda3/envs/cueq/lib/python3.10/site-packages/mace/tools/scripts_utils.py", line 184, in extract_config_mace_model
    mlp_irreps = o3.Irreps(f"{model_mlp_irreps.count((0, 1)) // len(heads)}x0e")
AttributeError: 'int' object has no attribute 'count'

I think it's due to the following code in mace/tools/scripts_utils.py:

    model_mlp_irreps = (
        o3.Irreps(str(model.readouts[-1].hidden_irreps))
        if model.num_interactions.item() > 1
        else 1
    )
    mlp_irreps = o3.Irreps(f"{model_mlp_irreps.count((0, 1)) // len(heads)}x0e")

If num_interactions=1, model_mlp_irreps=1 which does not have the .count method.

[ilyes319]

ah yeah I totally did not have a one layer model in mind, I will solve that but that might take me a decent amount of time to get to it.

[lucasdekam]

No problem, it's not urgent (with enable_cueq=False training still works).