feat: add output

deepmodeling · Mar 6, 2024 · 2da5a35 · 2da5a35
1 parent 1ae79a9
commit 2da5a35
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py
@@ -198,7 +198,7 @@ def compute_output_stats(
                     sampled = merged
 
                 sys_matrix, polar_bias = [], []
-                for sys in len(sampled):
+                for sys in range(len(sampled)):
                     if sampled[sys]["find_atomic_polarizability"] > 0.0:
                         for itype in range(self.ntypes):
                             # this is a tensor of shape nframes, nall
@@ -275,5 +275,7 @@ def forward(
             "bim,bmj->bij", gr.transpose(1, 2), out
         )  # (nframes * nloc, 3, 3)
         out = out.view(nframes, nloc, 3, 3)
+        if self.shift_diag:
+            out = out + self.constant_matrix[atype]*torch.eye(3, device=env.DEVICE)* self.scale[atype]
 
         return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}