Adding batch per epoch option to the finetune script (#10)

* Adding batch per epoch option to the finetune script Allow to specify the batch per epoch in the finetune script. * Remove copy paste error :)
DeepGraphLearning · Jan 13, 2024 · a1a9caf · a1a9caf
1 parent 33c6e6b
commit a1a9caf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/script/run.py b/script/run.py
@@ -289,12 +289,12 @@ def test(cfg, model, test_data, device, logger, filtered_data=None, return_metri
     val_filtered_data = val_filtered_data.to(device)
     test_filtered_data = test_filtered_data.to(device)
 
-    train_and_validate(cfg, model, train_data, valid_data, filtered_data=val_filtered_data, device=device, logger=logger)
+    train_and_validate(cfg, model, train_data, valid_data, filtered_data=val_filtered_data, device=device, batch_per_epoch=cfg.train.batch_per_epoch, logger=logger)
     if util.get_rank() == 0:
         logger.warning(separator)
         logger.warning("Evaluate on valid")
     test(cfg, model, valid_data, filtered_data=val_filtered_data, device=device, logger=logger)
     if util.get_rank() == 0:
         logger.warning(separator)
         logger.warning("Evaluate on test")
-    test(cfg, model, test_data, filtered_data=test_filtered_data, device=device, logger=logger)
+    test(cfg, model, test_data, filtered_data=test_filtered_data, device=device, logger=logger)