Wrote unit test for to_device
@@ -210,4 +210,54 @@ class TestConcatExamplesWithPadding(unittest.TestCase):
self.check_concat_dicts_padding(cuda.cupy)
+@testing.parameterize(
+ {'device': None, 'src_xp': numpy, 'dst_xp': numpy},
+ {'device': -1, 'src_xp': numpy, 'dst_xp': numpy},
+)
+class TestToDeviceCPU(unittest.TestCase):
+
+ def test_to_device(self):
+ x = self.src_xp.array([1], 'i')
+ y = dataset.to_device(self.device, x)
+ self.assertIsInstance(y, self.dst_xp.ndarray)
+
+
+@testing.parameterize(
+ {'device': None, 'src_xp': cuda.cupy, 'dst_xp': cuda.cupy},
+
+ {'device': -1, 'src_xp': cuda.cupy, 'dst_xp': numpy},
+
+ {'device': 0, 'src_xp': numpy, 'dst_xp': cuda.cupy},
+ {'device': 0, 'src_xp': cuda.cupy, 'dst_xp': cuda.cupy},
+)
+class TestToDeviceGPU(unittest.TestCase):
+
+ @attr.gpu
+ def test_to_device(self):
+ x = self.src_xp.array([1], 'i')
+ y = dataset.to_device(self.device, x)
+ self.assertIsInstance(y, self.dst_xp.ndarray)
+
+ if self.device >= 0:
+ self.assertEqual(int(y.device), self.device)
+
+ if self.device is None and self.src_xp is cuda.cupy:
+ self.assertEqual(int(x.device), int(y.device))
+
+
+@testing.parameterize(
+ {'device': 1, 'src_xp': numpy, 'dst_xp': cuda.cupy},
+ {'device': 1, 'src_xp': cuda.cupy, 'dst_xp': cuda.cupy},
+)
+class TestToDeviceMultiGPU(unittest.TestCase):
+
+ @attr.multi_gpu(2)
+ def test_to_device(self):
+ x = self.src_xp.array([1], 'i')
+ y = dataset.to_device(self.device, x)
+ self.assertIsInstance(y, self.dst_xp.ndarray)
+
+ self.assertEqual(int(y.device), self.device)
+
+
testing.run_module(__name__, __file__)
