Device error when evaluating MACE.

[jungsdao]

Hello,
First of all, thanks for making this wonderful tool public :)
I wanted to evaluate some structures with fitted MACE with eval_configs.py script.
This was trained with GPU.
While I was trying to evaluate some structures, I've received following error.

(/u/hjung/conda-envs/mace_env) hjung@raven02:~/4_Free_energy_sampling/1_MACE_test> python ~/Softwares/mace/scripts/eval_configs.py --configs="new_training_set_iter7.xyz" --model="checkpoints/MACE_model_run-123.model" --output="./output.xyz" --device="cpu"
Traceback (most recent call last):
  File "/u/hjung/Softwares/mace/scripts/eval_configs.py", line 151, in <module>
    main()
  File "/u/hjung/Softwares/mace/scripts/eval_configs.py", line 82, in main
    model = torch.load(f=args.model, map_location=args.device)
  File "/u/hjung/conda-envs/mace_env/lib/python3.9/site-packages/torch/serialization.py", line 809, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/u/hjung/conda-envs/mace_env/lib/python3.9/site-packages/torch/serialization.py", line 1172, in _load
    result = unpickler.load()
  File "/u/hjung/conda-envs/mace_env/lib/python3.9/site-packages/e3nn/util/codegen/_mixin.py", line 109, in __setstate__
    smod = torch.jit.load(buffer)
  File "/u/hjung/conda-envs/mace_env/lib/python3.9/site-packages/torch/jit/_serialization.py", line 164, in load
    cpp_module = torch._C.import_ir_module_from_buffer(
RuntimeError: No CUDA GPUs are available

Having seen complaints on CUDA GPU, I submitted evaluation job with the same configuration as in training with GPU and it succeeded.
Is it true that if I have trained a MACE model with GPU, then evaluation also has to be carried out in GPU? (i.e. If it's trained with GPU, it cannot be evaluated with CPU?)
I just want to be clarified with this issue.
Many thanks in advance!

Best regards,
Hyunwook

[ilyes319]

Hey @jungsdao, thank you very much for using MACE.

Due to a (probable) torchscript bug, a model needs to be saved on the CPU to be loaded to the CPU.
If you want to do this, you can add a --save_cpu flag to your input script, and it should work fine.
You can restart the train for like 0 epoch for your already trained model, just adding this flag. It will create a suitable model.
Or you can do the following trick on a machine with a GPU:

model = torch.load(path_to_model)
model_cpu = model.to("cpu")
torch.save(model_cpu, path_to_model_cpu)

Then you can use the model_cpu for your CPU evals.

Best,

Ilyes