# CUDA_VISIBLE_DEVICES=9 python3 mace/cli/run_train.py --name="MACE_model" --train_file="datasets/3BPA/train_300K.xyz" --valid_fraction=0.05 --test_file="datasets/3BPA/test_600K.xyz" --config_type_weights='{"Default":1.0}' --E0s='{1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639}' --model="MACE" --num_channels=256 --max_L=3 --r_max=5.0 --batch_size=16 --max_num_epochs=1500 --stage_two --start_stage_two=1200 --ema --ema_decay=0.99 --amsgrad --restart_latest --device=cuda 


CUDA_VISIBLE_DEVICES=9 python3 ./scripts/run_train.py \
  --name="MACE_model" \
  --train_file="datasets/chignolin_data/xyz/train.xyz" \
  --valid_fraction=0.05 \
  --forces_key="forces" \
  --energy_key="energy" \
  --test_file="datasets/chignolin_data/xyz/test.xyz" \
  --E0s='{1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639, 16: -10950.6658371305}' \
  --atomic_numbers='[1, 6, 7, 8, 16]' \
  --model="ScaleShiftMACE" \
  --num_channels=64 \
  --max_L=3 \
  --r_max=9.0 \
  --batch_size=2 \
  --max_num_epochs=100 \
  --ema \
  --ema_decay=0.99 \
  --amsgrad \
  --default_dtype="float32" \
  --device=cuda \
  --seed=123 \
  --swa