Parallelisation

[yangjackie]

Hi,

Thanks for sharing this code. I had a little play with it on a garnet material and it works fairly good on training sets include finite temperature displaced structure generated from Phonopy.

So far I only managed to run this on a singe core. Just wondering how can I run the training on parallel architecture? Can you provide an example how to set it up? I just dive into N.N. recently, so not quite familiar with the JAX and other libraries that are used in your code.

Thanks a lot Jack