python ./mace/mace/cli/run_train.py \
    --name="mace-omat-medium" \
    --config=config-omat.yaml \
    --train_file="/train/all_data" \
    --valid_file="/val/all_data" \
    --test_dir="/val/all_data" \
    --statistics_file="omat-statistics.json" \
    --interaction_first="RealAgnosticDensityInteractionBlock" \
    --interaction="RealAgnosticDensityResidualInteractionBlock" \
    --energy_weight=1 \
    --forces_weight=10 \
    --compute_stress=True \
    --stress_weight=10 \
    --stress_key='stress' \
    --energy_key='energy' \
    --forces_key='forces' \
    --eval_interval=1 \
    --error_table='PerAtomMAE' \
    --model="MACE" \
    --hidden_irreps="128x0e + 128x1o" \
    --mean=0.0 \
    --std=1.0 \
    --atomic_numbers="[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 89, 90, 91, 92, 93, 94]" \
    --num_interactions=2 \
    --r_max=6.0 \
    --eval_interval=1 \
    --batch_size=16 \
    --valid_batch_size=16 \
    --pair_repulsion \
    --distance_transform="Agnesi" \
    --ema \
    --ema_decay=0.999 \
    --error_table="PerAtomMAE" \
    --max_num_epochs=10 \
    --device=cuda \
    --restart_latest \
    --weight_decay=1e-8 \
    --save_cpu \
    --default_dtype="float64" \
    --seed=1 \
    --clip_grad=10.0 \
    --num_workers=2 \
    --distributed \
    --launcher="slurm" \
    --results_dir="./results_omat" \
    --log_dir="./logs_omat"