fix: use get_model_def_script method (deepmodeling#3843)

njzjz · May 31, 2024 · d4b6c41 · d4b6c41
1 parent e3acea5
commit d4b6c41
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
@@ -303,7 +303,7 @@ def train(FLAGS):
 
 def freeze(FLAGS):
     model = torch.jit.script(inference.Tester(FLAGS.model, head=FLAGS.head).model)
-    if '"type": "dpa2"' in model.model_def_script:
+    if '"type": "dpa2"' in model.get_model_def_script():
         extra_files = {"type": "dpa2"}
     else:
         extra_files = {"type": "else"}