Update output_def.py

deepmodeling · Mar 2, 2024 · 967736a · 967736a
1 parent f3c7bf2
commit 967736a
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import functools
-import math
 from enum import (
     IntEnum,
 )
@@ -194,7 +193,11 @@ def __init__(
     ):
         self.name = name
         self.shape = list(shape)
-        self.output_size = math.prod(self.shape)
+        # jit doesn't support math.prod(self.shape)
+        self.output_size = 1
+        len_shape = len(self.shape)
+        for i in range(len_shape):
+            self.output_size *= self.shape[i]
         self.atomic = atomic
         self.reduciable = reduciable
         self.r_differentiable = r_differentiable