diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py
index 16c097f743..26e3fcd266 100644
--- a/deepmd/entrypoints/test.py
+++ b/deepmd/entrypoints/test.py
@@ -1277,6 +1277,7 @@ def test_dipole(
         high_prec=False,
         type_sel=dp.get_sel_type(),
     )
+
     test_data = data.get_test()
     dipole, numb_test, atype = run_test(dp, test_data, numb_test, data)
 
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 26a27c82d7..df11ec4772 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -750,7 +750,9 @@ def _load_data(
             data = path.load_numpy().astype(dtype)
             try:  # YWolfeee: deal with data shape error
                 if atomic:
-                    if type_sel is not None:
+                    if (
+                        type_sel is not None and natoms_sel != natoms
+                    ):  # handle PT all atom case
                         # check the data shape is nsel or natoms
                         if data.size == nframes * natoms_sel * ndof_:
                             if output_natoms_for_type_sel: