fix auto pin_memory : idist.device().type should be used (pytorch#1131)

* fix auto pin_memory : idist.device().type should be used * fix cuda in device * fix test * use idist.device().type to test * add missing () Co-authored-by: Desroziers <sylvain.desroziers@ifpen.fr>
geotyper · Jun 15, 2020 · db63e94 · db63e94
1 parent 3b1dc3a
commit db63e94
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
@@ -98,7 +98,7 @@ def auto_dataloader(dataset, **kwargs):
         )
         kwargs["pin_memory"] = False
     else:
-        kwargs["pin_memory"] = kwargs.get("pin_memory", idist.device() == "cuda")
+        kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type)
 
     logger.info("Use data loader kwargs for dataset '{}': \n\t{}".format(repr(dataset)[:20].strip(), kwargs))
     dataloader = DataLoader(dataset, **kwargs)

diff --git a/tests/ignite/distributed/test_auto.py b/tests/ignite/distributed/test_auto.py
@@ -40,7 +40,7 @@ def _test_auto_dataloader(ws, nproc, sampler_name=None, dl_type=DataLoader):
         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")
+        assert dataloader.pin_memory == ("cuda" in idist.device().type)
 
 
 def _test_auto_model_optimizer(ws, device):