Update stage3.py

lsh0520 · Apr 2, 2024 · db880a3 · db880a3
1 parent 872993c
commit db880a3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stage3.py b/stage3.py
@@ -76,7 +76,7 @@ def main(args):
         logger=logger,
     )
 
-    if args.mode.find('pretrain') >= 0:
+    if args.mode.find('train') >= 0:
         trainer.fit(model, datamodule=dm)
         trainer.test(model, datamodule=dm)
     elif args.mode.find('eval') >= 0: