[GraphBolt] avoid warning of pickle local function (#6284)

python/dgl/graphbolt/item_sampler.py

@@ -282,24 +282,9 @@ class ItemSampler(IterDataPipe):
         self._drop_last = drop_last
         self._shuffle = shuffle
 
-    def __iter__(self) -> Iterator:
-        data_pipe = IterableWrapper(self._item_set)
-        # Shuffle before batch.
-        if self._shuffle:
-            # `torchdata.datapipes.iter.Shuffler` works with stream too.
-            # To ensure randomness, make sure the buffer size is at least 10
-            # times the batch size.
-            buffer_size = max(10000, 10 * self._batch_size)
-            data_pipe = data_pipe.shuffle(buffer_size=buffer_size)
-
-        # Batch.
-        data_pipe = data_pipe.batch(
-            batch_size=self._batch_size,
-            drop_last=self._drop_last,
-        )
-
-        # Collate.
+    @staticmethod
     def _collate(batch):
+        """Collate items into a batch. For internal use only."""
         data = next(iter(batch))
         if isinstance(data, DGLGraph):
             return dgl_batch(batch)
@@ -316,7 +301,24 @@ class ItemSampler(IterDataPipe):
             }
         return default_collate(batch)
 
-        data_pipe = data_pipe.collate(collate_fn=partial(_collate))
+    def __iter__(self) -> Iterator:
+        data_pipe = IterableWrapper(self._item_set)
+        # Shuffle before batch.
+        if self._shuffle:
+            # `torchdata.datapipes.iter.Shuffler` works with stream too.
+            # To ensure randomness, make sure the buffer size is at least 10
+            # times the batch size.
+            buffer_size = max(10000, 10 * self._batch_size)
+            data_pipe = data_pipe.shuffle(buffer_size=buffer_size)
+
+        # Batch.
+        data_pipe = data_pipe.batch(
+            batch_size=self._batch_size,
+            drop_last=self._drop_last,
+        )
+
+        # Collate.
+        data_pipe = data_pipe.collate(collate_fn=self._collate)
 
         # Map to minibatch.
         data_pipe = data_pipe.map(