Use batch_per_epoch to slice the train data (#11)

DeepGraphLearning · Jan 13, 2024 · 04d5c13 · 04d5c13
1 parent a1a9caf
commit 04d5c13
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/script/run.py b/script/run.py
@@ -2,6 +2,7 @@
 import sys
 import math
 import pprint
+from itertools import islice
 
 import torch
 import torch_geometric as pyg
@@ -59,7 +60,7 @@ def train_and_validate(cfg, model, train_data, valid_data, device, logger, filte
 
             losses = []
             sampler.set_epoch(epoch)
-            for batch in train_loader:
+            for batch in islice(train_loader, batch_per_epoch):
                 batch = tasks.negative_sampling(train_data, batch, cfg.task.num_negative,
                                                 strict=cfg.task.strict_negative)
                 pred = parallel_model(train_data, batch)