fix: use get_model_def_script method

deepmodeling · May 31, 2024 · ae0fa4e · ae0fa4e
1 parent 3a7fbcf
commit ae0fa4e
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"}