typo

pyg-team · Jan 19, 2023 · 5c6718e · 5c6718e
1 parent b9b5b51
commit 5c6718e
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/torch_geometric/data/lightning/datamodule.py b/torch_geometric/data/lightning/datamodule.py
@@ -41,17 +41,16 @@ def __init__(self, has_val: bool, has_test: bool, **kwargs):
 
         kwargs.setdefault('batch_size', 1)
         kwargs.setdefault('num_workers', 0)
+        kwargs.setdefault('pin_memory', True)
+        kwargs.setdefault('persistent_workers',
+                          kwargs.get('num_workers', 0) > 0)
 
         if 'shuffle' in kwargs:
             warnings.warn(f"The 'shuffle={kwargs['shuffle']}' option is "
                           f"ignored in '{self.__class__.__name__}'. Remove it "
                           f"from the argument list to disable this warning")
             del kwargs['shuffle']
 
-        kwargs.setdefault('pin_memory', True)
-        kwargs.setdefault('persistent_workers',
-                          kwargs.get('num_workers', 0) > 0)
-
         self.kwargs = kwargs
 
     def prepare_data(self):