From cc67fdd1307e6b95e3bf26094386a99071b790a5 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 7 Nov 2024 04:45:37 -0500
Subject: [PATCH] fix type_embed.py

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/tf/utils/type_embed.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deepmd/tf/utils/type_embed.py b/deepmd/tf/utils/type_embed.py
index 9fae1196dc..9b7b17528d 100644
--- a/deepmd/tf/utils/type_embed.py
+++ b/deepmd/tf/utils/type_embed.py
@@ -6,6 +6,8 @@
     Union,
 )
 
+import numpy as np
+
 from deepmd.dpmodel.utils.network import (
     EmbeddingNet,
 )
@@ -327,7 +329,7 @@ def serialize(self, suffix: str = "") -> dict:
             layer_idx = int(m[1]) - 1
             weight_name = m[0]
             if weight_name == "idt":
-                if value == 0.0:
+                if not isinstance(value, np.ndarray):
                     # ignore 0.0 set by deserialize
                     continue
                 value = value.ravel()