Restore support for passing device in CPU-only environment (#107)

Add test for passing device to check this works on CPU
NVIDIA-Merlin · Mar 13, 2023 · 226ad69 · 226ad69
1 parent 3c11b59
commit 226ad69
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/merlin/dataloader/loader_base.py b/merlin/dataloader/loader_base.py
@@ -74,10 +74,8 @@ def __init__(
         self.global_rank = global_rank or 0
         self.drop_last = drop_last
 
-        if device:
-            self.device = device
-        else:
-            self.device = "cpu" if not HAS_GPU or dataset.cpu else 0
+        device = device or 0
+        self.device = "cpu" if not HAS_GPU or dataset.cpu else device
 
         if self.device == "cpu":
             self._array_lib = np

diff --git a/tests/unit/dataloader/test_tf_dataloader.py b/tests/unit/dataloader/test_tf_dataloader.py
@@ -98,6 +98,11 @@ def test_simple_model():
     _ = model.evaluate(loader)
 
 
+def test_with_device():
+    dataset = Dataset(make_df({"a": [1]}))
+    tf_dataloader.Loader(dataset, batch_size=1, device=1).peek()
+
+
 def test_nested_list():
     num_rows = 100
     batch_size = 12