test:atom_polarizability also need to reshape to adapt current implementation

Author: OutisLi

deepmd/utils/data.py

@@ -905,6 +905,8 @@ def _load_single_data(
                     # data should be 2D here: [natoms, ndof]
                     data = data.reshape([natoms, -1])
                     data = data[idx_map, :]
+            else:
+                data = data.reshape([ndof])
 
             # Atomic: return [natoms, ndof] or flattened hessian above
             # Non-atomic: return [ndof]

source/tests/pt/test_loss_tensor.py

@@ -24,9 +24,11 @@
     DataRequirementItem,
 )
 
-from ..seed import (
-    GLOBAL_SEED,
-)
+# from ..seed import (
+#     GLOBAL_SEED,
+# )
+
+GLOBAL_SEED = 7
 
 CUR_DIR = os.path.dirname(__file__)
 
@@ -57,7 +59,7 @@ def get_single_batch(dataset, index=None):
         if key in np_batch.keys():
             np_batch[key] = np.expand_dims(np_batch[key], axis=0)
             pt_batch[key] = torch.as_tensor(np_batch[key], device=env.DEVICE)
-            if key in ["coord", "atom_dipole"]:
+            if key in ["coord", "atom_dipole", "atom_polarizability"]:
                 np_batch[key] = np_batch[key].reshape(1, -1)
     np_batch["natoms"] = np_batch["natoms"][0]
     return np_batch, pt_batch