Auto pin_memory

examples/contrib/cifar10/main.py

@@ -192,12 +192,7 @@ def get_dataflow(config):
 
     # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
     train_loader = idist.auto_dataloader(
-        train_dataset,
-        batch_size=config["batch_size"],
-        num_workers=config["num_workers"],
-        shuffle=True,
-        pin_memory="cuda" in idist.device().type,
-        drop_last=True,
+        train_dataset, batch_size=config["batch_size"], num_workers=config["num_workers"], shuffle=True, drop_last=True,
     )
 
     test_loader = idist.auto_dataloader(

ignite/distributed/auto.py

@@ -25,6 +25,7 @@ def auto_dataloader(dataset, **kwargs):
     - number of workers is scaled by number of local processes: ``num_workers / nprocs``.
     - if no sampler provided by user, `torch DistributedSampler` is setup.
     - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
+    - if the default device is 'cuda', `pin_memory` is automatically set to `True`.
 
     .. warning::
 
@@ -96,6 +97,8 @@ def auto_dataloader(dataset, **kwargs):
             "Argument `pin_memory=False` will be used to construct data loader."
         )
         kwargs["pin_memory"] = False
+    else:
+        kwargs["pin_memory"] = kwargs.get("pin_memory", idist.device() == "cuda")
 
     logger.info("Use data loader kwargs for dataset '{}': \n\t{}".format(repr(dataset)[:20].strip(), kwargs))
     dataloader = DataLoader(dataset, **kwargs)

tests/ignite/distributed/test_auto.py

@@ -39,6 +39,8 @@ def _test_auto_dataloader(ws, nproc, sampler_name=None, dl_type=DataLoader):
     else:
         sampler_type = DistributedSampler if sampler is None else DistributedProxySampler
         assert isinstance(dataloader.sampler, sampler_type)
+    if isinstance(dataloader, DataLoader):
+        assert dataloader.pin_memory == (data.device == "cuda")
 
 
 def _test_auto_model_optimizer(ws, device):