From cdb1d3dd96a5fb72eb94ee49ec706864bc849a8e Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Sat, 3 Aug 2024 06:27:01 +0200
Subject: [PATCH] Upgrade to `metatensor-torch` 0.5.3

---
 pyproject.toml             | 2 +-
 src/metatrain/cli/train.py | 3 ---
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index dfdb38595..2c1d7c747 100644
--- 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
index 11026e265..e567b5c87 100644
--- 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()