use math.prod

wanghan-iapcm · Feb 13, 2024 · 6573e19 · 6573e19
1 parent f401eb6
commit 6573e19
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
+import math
 from typing import (
     Dict,
     List,
@@ -139,9 +140,7 @@ def _cal_hessian_all(
             for kk in hess_keys:
                 vdef = fdef[kk]
                 vshape = vdef.shape
-                vsize = 1
-                for ii in vshape:
-                    vsize *= ii
+                vsize = math.prod(vdef.shape)
                 # loop over frames
                 for ii in range(nf):
                     icoord = coord[ii]