[GraphBolt][CUDA] Fix link prediction early-stop. (#7083)

examples/sampling/graphbolt/link_prediction.py

@@ -332,6 +332,8 @@ def train(args, model, graph, features, train_set):
 
             total_loss += loss.item()
             if step + 1 == args.early_stop:
+                # Early stopping requires a new dataloader to reset its state.
+                dataloader = create_dataloader(args, graph, features, train_set)
                 break
 
         end_epoch_time = time.time()

python/dgl/graphbolt/feature_fetcher.py

@@ -174,10 +174,5 @@ class FeatureFetcher(MiniBatchTransformer):
         with torch.cuda.stream(self.stream):
             data = self._read_data(data, current_stream)
             if self.stream is not None:
-                event = torch.cuda.current_stream().record_event()
-
-                def _wait():
-                    event.wait()
-
-                data.wait = _wait
+                data.wait = torch.cuda.current_stream().record_event().wait
             return data