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

Add test for passing device to check this works on CPU

Author: oliverholworthy

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

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