openkim / kliff

KIM-based Learning-Integrated Fitting Framework for interatomic potentials.
https://kliff.readthedocs.io
GNU Lesser General Public License v2.1
34 stars 20 forks source link

Add support to train NN model on GPU #23

Closed mjwen closed 2 years ago

mjwen commented 2 years ago

To train on GPU, use the gpu param, which can be an integer (which gives the index of the gpu) and can be True (will use the first gpu). For example, CalculatorTorch(model, gpu=0) and CalculatorTorch(model, gpu=True).

Same for multiple species calculator CalculatorTorchSeparateSpecies.