Update VE.py

SuJQ0722 · Oct 26, 2021 · 9e9a5e9 · 9e9a5e9
1 parent 69e9225
commit 9e9a5e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/VE.py b/VE.py
@@ -213,6 +213,7 @@ def main(args, config):
                     }
                     torch.save(save_obj, os.path.join(args.output_dir, 'checkpoint_best.pth')) 
                     best = float(val_stats['acc'])
+                    best_epoch = epoch
 
         if args.evaluate:
             break
@@ -248,4 +249,4 @@ def main(args, config):
 
     yaml.dump(config, open(os.path.join(args.output_dir, 'config.yaml'), 'w'))    
 
-    main(args, config)
+    main(args, config)