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()