diff --git a/pyproject.toml b/pyproject.toml
index e80ed1c1..07a75e75 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -64,7 +64,7 @@ alchemical-model = [
   "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@51ff519",
 ]
 pet = [
-  "pet @ git+https://github.com/lab-cosmo/pet@292ab40",
+  "pet @ git+https://github.com/lab-cosmo/pet@8841021",
 ]
 gap = [
     "rascaline-torch @ git+https://github.com/luthaf/rascaline@d181b28#subdirectory=python/rascaline-torch",
diff --git a/src/metatrain/experimental/pet/default-hypers.yaml b/src/metatrain/experimental/pet/default-hypers.yaml
index 422a5493..9b283c75 100644
--- a/src/metatrain/experimental/pet/default-hypers.yaml
+++ b/src/metatrain/experimental/pet/default-hypers.yaml
@@ -35,6 +35,8 @@ model:
 
 training:
   INITIAL_LR: 1e-4
+  EPOCH_NUM: 1000
+  EPOCHS_WARMUP: 50
   EPOCH_NUM_ATOMIC: 1000000000000000000 # structural version is called "EPOCH_NUM"
   SCHEDULER_STEP_SIZE_ATOMIC: 500000000 # structural version is called "SCHEDULER_STEP_SIZE"
   EPOCHS_WARMUP_ATOMIC: 250000000 # structural version is called "EPOCHS_WARMUP"
diff --git a/tox.ini b/tox.ini
index d388b99e..77463661 100644
--- a/tox.ini
+++ b/tox.ini
@@ -116,7 +116,7 @@ extras = pet
 changedir = src/metatrain/experimental/pet/tests/
 commands =
     python -m pip uninstall -y pet
-    python -m pip install "pet @ git+https://github.com/lab-cosmo/pet@292ab40" --no-build-isolation
+    python -m pip install "pet @ git+https://github.com/lab-cosmo/pet@8841021" --no-build-isolation
     pytest {posargs}
 
 [testenv:gap-tests]