Update loss.py

align unit
deepmodeling · Dec 31, 2024 · 24c6afe · 24c6afe
1 parent 46c6aa7
commit 24c6afe
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dptb/nnops/loss.py b/dptb/nnops/loss.py
@@ -323,6 +323,7 @@ def __init__(
         self.device = device
         self.onsite_shift = onsite_shift
         self.coeff_ham = coeff_ham
+        assert self.coeff_ham <= 1.
         self.coeff_ovp = coeff_ovp
 
         if basis is not None:
@@ -393,7 +394,7 @@ def forward(self, data: AtomicDataDict, ref_data: AtomicDataDict):
 
         eigloss = self.eigloss(data, ref_data)
 
-        return self.coeff_ham * ham_loss + eigloss
+        return self.coeff_ham * ham_loss + (1 - self.coeff_ham) eigloss
 
 
 
@@ -1106,4 +1107,4 @@ def __cal_norm__(self, irreps: Irreps, x: torch.Tensor):
             tensor = tensor.norm(dim=-1)
             out.append(tensor)
 
-        return torch.cat(out, dim=-1).squeeze(0)
+        return torch.cat(out, dim=-1).squeeze(0)