diff --git a/src/aiida_quantumespresso_hp/utils/general.py b/src/aiida_quantumespresso_hp/utils/general.py
index 3566bac..99c9fd0 100644
--- a/src/aiida_quantumespresso_hp/utils/general.py
+++ b/src/aiida_quantumespresso_hp/utils/general.py
@@ -28,7 +28,7 @@ def is_perturb_only_atom(parameters: dict) -> int | None:
     match = None  # making sure that if the dictionary is empty we don't raise an `UnboundLocalError`
 
     for key in parameters.keys():
-        match = re.search(r'perturb_only_atom.*([0-9]).*', key)
+        match = re.search(r'perturb_only_atom.*?(\d+).*', key)
         if match:
             if not parameters[key]:  # also the key must be `True`
                 match = None  # making sure to have `None`
diff --git a/tests/utils/test_general.py b/tests/utils/test_general.py
index 031947f..c7e3567 100644
--- a/tests/utils/test_general.py
+++ b/tests/utils/test_general.py
@@ -26,5 +26,8 @@ def test_is_perturb_only_atom():
     parameters = {'perturb_only_atom(1)': True}
     assert is_perturb_only_atom(parameters) == 1
 
+    parameters = {'perturb_only_atom(20)': True}
+    assert is_perturb_only_atom(parameters) == 20
+
     parameters = {'perturb_only_atom(1)': False}
     assert is_perturb_only_atom(parameters) is None