pt: cleanup tester

deepmodeling · Mar 9, 2024 · bb324ff · bb324ff
1 parent a9bcf41
commit bb324ff
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 354 deletions.
diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
@@ -281,9 +281,7 @@ def train(FLAGS):
 
 
 def freeze(FLAGS):
-    model = torch.jit.script(
-        inference.Tester(FLAGS.model, numb_test=1, head=FLAGS.head).model
-    )
+    model = torch.jit.script(inference.Tester(FLAGS.model, head=FLAGS.head).model)
     torch.jit.save(
         model,
         FLAGS.output,