From 00105c7d548e297724d223d5667245d0e84102bf Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Wed, 28 Feb 2024 20:47:21 +0800
Subject: [PATCH] Update make_base_descriptor.py

---
 deepmd/dpmodel/descriptor/make_base_descriptor.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
index fe911551d5..0f55916111 100644
--- a/deepmd/dpmodel/descriptor/make_base_descriptor.py
+++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -4,8 +4,10 @@
     abstractmethod,
 )
 from typing import (
+    Callable,
     List,
     Optional,
+    Union,
 )
 
 from deepmd.common import (
@@ -84,7 +86,9 @@ def mixed_types(self) -> bool:
             """
             pass
 
-        def compute_input_stats(self, merged: callable, path: Optional[DPPath] = None):
+        def compute_input_stats(
+            self, merged: Union[Callable, List], path: Optional[DPPath] = None
+        ):
             """Update mean and stddev for descriptor elements."""
             raise NotImplementedError