From 967736a526eb48e44e6159dd4afeeae02c8d3bf1 Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Sun, 3 Mar 2024 02:43:06 +0800
Subject: [PATCH] Update output_def.py

---
 deepmd/dpmodel/output_def.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py
index 0bc248495a..ac41513246 100644
--- 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