[GraphBolt][CUDA] Fix link prediction early-stop.

Author: mfbalin

examples/sampling/graphbolt/link_prediction.py

@@ -340,6 +340,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 @@ def _read(self, data):
         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