Update repformer_layer.py

deepmodeling · May 9, 2024 · e329e30 · e329e30
1 parent a8c89dc
commit e329e30
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py
@@ -834,7 +834,7 @@ def _cal_hg(
             invnnei = (1.0 / float(nnei)) * torch.ones(
                 (nb, nloc, 1, 1), dtype=g2.dtype, device=g2.device
             )
-        # nb x nloc x 3 x ng
+        # nb x nloc x 3 x ng2
         h2g2 = torch.matmul(torch.transpose(h2, -1, -2), g2) * invnnei
         return h2g2