Upgrade to metatensor-torch 0.5.3 (#328)

metatensor · Aug 3, 2024 · 6152ba6 · 6152ba6
1 parent 96f2c74
commit 6152ba6
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "ase < 3.23.0",
     "metatensor-learn==0.2.2",
     "metatensor-operations==0.2.1",
-    "metatensor-torch==0.5.2",
+    "metatensor-torch==0.5.3",
     "jsonschema",
     "omegaconf",
     "python-hostlist",

diff --git a/src/metatrain/cli/train.py b/src/metatrain/cli/train.py
@@ -382,9 +382,6 @@ def train_model(
             mts_atomistic_model.buffers(),
         )
     ).device
-    # always save on CPU. TODO: remove after release of
-    # https://github.com/lab-cosmo/metatensor/pull/668
-    mts_atomistic_model = mts_atomistic_model.to("cpu")
     mts_atomistic_model.save(str(output_checked), collect_extensions=extensions_path)
     # the model is first saved and then reloaded 1) for good practice and 2) because
     # MetatensorAtomisticModel only torchscripts (makes faster) during save()