From 977447b77e80c096667263cd967868427614fd9b Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 26 Nov 2024 14:24:33 +0800
Subject: [PATCH] fix: add model_def_script to zbl

---
 deepmd/pt/model/model/__init__.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py
index da1edb3354..9d4bf45f15 100644
--- a/deepmd/pt/model/model/__init__.py
+++ b/deepmd/pt/model/model/__init__.py
@@ -196,7 +196,7 @@ def get_zbl_model(model_params):
     rmax = model_params["sw_rmax"]
     atom_exclude_types = model_params.get("atom_exclude_types", [])
     pair_exclude_types = model_params.get("pair_exclude_types", [])
-    return DPZBLModel(
+    model = DPZBLModel(
         dp_model,
         pt_model,
         rmin,
@@ -205,6 +205,8 @@ def get_zbl_model(model_params):
         atom_exclude_types=atom_exclude_types,
         pair_exclude_types=pair_exclude_types,
     )
+    model.model_def_script = json.dumps(model_params)
+    return model
 
 
 def _can_be_converted_to_float(value) -> Optional[bool]: