Fix: Pt Polar Stat keyword

deepmd/pt/model/task/polarizability.py

@@ -230,12 +230,12 @@
 
                 sys_constant_matrix = []
                 for sys in range(len(sampled)):
-                    nframs = sampled[sys]["type"].shape[0]
+                    nframs = sampled[sys]["atype"].shape[0]
 
                     if sampled[sys]["find_atomic_polarizability"] > 0.0:
                         sys_atom_polar = compute_stats_from_atomic(
                             sampled[sys]["atomic_polarizability"].numpy(force=True),
-                            sampled[sys]["type"].numpy(force=True),
+                            sampled[sys]["atype"].numpy(force=True),
                         )[0]
                     else:
                         if not sampled[sys]["find_polarizability"] > 0.0:
@@ -244,7 +244,7 @@
                             (nframs, self.ntypes), dtype=env.GLOBAL_NP_FLOAT_PRECISION
                         )
                         for itype in range(self.ntypes):
-                            type_mask = sampled[sys]["type"] == itype
+                            type_mask = sampled[sys]["atype"] == itype
                             sys_type_count[:, itype] = type_mask.sum(dim=1).numpy(
                                 force=True
                             )

deepmd/pt/utils/stat.py

@@ -61,12 +61,17 @@
                         if dd not in sys_stat:
                             sys_stat[dd] = []
                         sys_stat[dd].append(stat_data[dd])
+                    elif isinstance(stat_data[dd], np.float32):
+                        sys_stat[dd] = stat_data[dd]
                     else:
                         pass
+
         for key in sys_stat:
-            if sys_stat[key] is None or sys_stat[key][0] is None:
+            if isinstance(sys_stat[key], np.float32):
+                pass
+            elif sys_stat[key] is None or sys_stat[key][0] is None:
                 sys_stat[key] = None
-            else:
+            elif isinstance(stat_data[dd], torch.Tensor):
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
         dict_to_device(sys_stat)
         lst.append(sys_stat)

deepmd/pt/utils/utils.py

@@ -116,6 +116,10 @@
     for key in sample_dict:
         if isinstance(sample_dict[key], list):
             sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]]
+        if isinstance(sample_dict[key], np.float32):
+            sample_dict[key] = (
+                torch.ones(1, dtype=torch.float32, device=DEVICE) * sample_dict[key]
+            )
         else:
             if sample_dict[key] is not None:
                 sample_dict[key] = sample_dict[key].to(DEVICE)

source/tests/pt/model/test_polar_stat.py

@@ -30,7 +30,7 @@ def setUp(self) -> None:
         find_atomic_polarizability = torch.rand(1, device=env.DEVICE)
         self.sampled = [
             {
-                "type": types,
+                "atype": types,
                 "find_atomic_polarizability": find_atomic_polarizability,
                 "atomic_polarizability": atomic_polarizability,
                 "polarizability": polarizability,
@@ -40,6 +40,7 @@ def setUp(self) -> None:
         self.all_stat = {
             k: [v.numpy(force=True)] for d in self.sampled for k, v in d.items()
         }
+        self.all_stat["type"] = self.all_stat.pop("atype")
         self.tfpolar = PolarFittingSeA(
             ntypes=ntypes,
             dim_descrpt=1,