| 393 | |
| 394 | @eval_modes() |
| 395 | def test_batch_to_gpu(): |
| 396 | input = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32) |
| 397 | t_cpu = ndd.tensor(input) |
| 398 | t_gpu = t_cpu.gpu() |
| 399 | assert t_gpu.device == ndd.Device("gpu") |
| 400 | b_gpu = ndd.Batch([t_gpu]) |
| 401 | b_gpu.evaluate() |
| 402 | assert b_gpu.device == ndd.Device("gpu") |
| 403 | assert b_gpu.dtype == ndd.int32 |
| 404 | assert b_gpu.batch_size == 1 |
| 405 | assert b_gpu.ndim == 2 |
| 406 | assert b_gpu.shape == [(2, 3)] |
| 407 | assert np.array_equal(asnumpy(b_gpu.tensors[0]), input) |
| 408 | |
| 409 | |
| 410 | @eval_modes() |