[DistDGL] back-compatible with dgl.dataloading (#7817)

dmlc · Oct 6, 2024 · 0dc0d57 · 0dc0d57
1 parent adbefc6
commit 0dc0d57
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
@@ -1502,6 +1502,16 @@ def __new__(cls, *args, **kwargs):
         return NewEdgeCollator(*args, **kwargs)
 
 
+def _remove_kwargs_dist(kwargs):
+    """Deprecated."""
+    if "num_workers" in kwargs:
+        del kwargs["num_workers"]
+    if "pin_memory" in kwargs:
+        del kwargs["pin_memory"]
+        print("Distributed DataLoaders do not support pin_memory.")
+    return kwargs
+
+
 class DistDataLoader:
     """Deprecated. Please use :class:`~dgl.distributed.DistDataLoader` instead."""