diff --git a/src/pao_io.F b/src/pao_io.F
index 2c6f4a0e52..534ff9e3e6 100644
--- a/src/pao_io.F
+++ b/src/pao_io.F
@@ -172,8 +172,9 @@ CONTAINS
 !> \param atom2kind ...
 !> \param positions ...
 !> \param xblocks ...
+!> \param ml_range ...
 ! **************************************************************************************************
-   SUBROUTINE pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks)
+   SUBROUTINE pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks, ml_range)
       CHARACTER(LEN=default_path_length), INTENT(IN)     :: filename
       CHARACTER(LEN=default_string_length), INTENT(OUT)  :: param
       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: hmat
@@ -181,6 +182,7 @@ CONTAINS
       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: atom2kind
       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: positions
       TYPE(pao_ioblock_type), ALLOCATABLE, DIMENSION(:)  :: xblocks
+      INTEGER, DIMENSION(2), INTENT(OUT), OPTIONAL       :: ml_range
 
       CHARACTER(len=*), PARAMETER :: routineN = 'pao_read_raw', routineP = moduleN//':'//routineN
 
@@ -266,6 +268,14 @@ CONTAINS
             READ (unit_nr, fmt=*) label, natoms
             ALLOCATE (positions(natoms, 3), atom2kind(natoms), xblocks(natoms))
             positions = 0.0_dp; atom2kind = -1
+            IF (PRESENT(ml_range)) ml_range = (/1, natoms/)
+
+         ELSE IF (TRIM(label) == "MLRange") THEN
+            ! Natoms entry has to come first
+            CPASSERT(natoms > 0)
+            ! range of atoms whose xblocks are used for machine learning
+            READ (unit_nr, fmt=*) label, i1, i2
+            IF (PRESENT(ml_range)) ml_range = (/i1, i2/)
 
          ELSE IF (TRIM(label) == "Atom") THEN
             READ (unit_nr, fmt=*) label, iatom, str_in, pos_in
diff --git a/src/pao_ml.F b/src/pao_ml.F
index ae2a77b8d9..f0c5175f7f 100644
--- a/src/pao_ml.F
+++ b/src/pao_ml.F
@@ -149,6 +149,7 @@ CONTAINS
       CHARACTER(LEN=default_string_length)               :: param
       INTEGER                                            :: iatom, ikind, natoms, nkinds, nparams
       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: atom2kind, kindsmap
+      INTEGER, DIMENSION(2)                              :: ml_range
       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: hmat, positions
       TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
       TYPE(cell_type), POINTER                           :: cell
@@ -167,7 +168,7 @@ CONTAINS
 
       ! parse training data on first rank
       IF (para_env%mepos == para_env%source) THEN
-         CALL pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks)
+         CALL pao_read_raw(filename, param, hmat, kinds, atom2kind, positions, xblocks, ml_range)
 
          ! check parametrization
          IF (TRIM(param) .NE. TRIM(ADJUSTL(id2str(pao%parameterization)))) &
@@ -192,6 +193,10 @@ CONTAINS
       CALL mp_bcast(kindsmap, para_env%source, para_env%group)
       CALL mp_bcast(atom2kind, para_env%source, para_env%group)
       CALL mp_bcast(positions, para_env%source, para_env%group)
+      CALL mp_bcast(ml_range, para_env%source, para_env%group)
+
+      IF (ml_range(1) /= 1 .OR. ml_range(2) /= natoms) &
+         CPWARN("Skipping some atoms for PAO-ML training.")
 
       ! create cell from read-in h-matrix
       CALL cell_create(cell, hmat)
@@ -210,6 +215,7 @@ CONTAINS
       ! however the input and output arrays will only be allocated on one rank per entry.
       ! We farm out the expensive calculation of the descriptor across ranks.
       DO iatom = 1, natoms
+         IF (iatom < ml_range(1) .OR. ml_range(2) < iatom) CYCLE
          ALLOCATE (new_point)
 
          ! training-point input, calculate descriptor only on one rank