pt: cleanup tester (#3442)

deepmodeling · Mar 10, 2024 · a286bd4 · a286bd4
1 parent fd82f04
commit a286bd4
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,