diff --git a/chainer_pytorch_migration/tensor.py b/chainer_pytorch_migration/tensor.py index 2438b53..56fa105 100644 --- a/chainer_pytorch_migration/tensor.py +++ b/chainer_pytorch_migration/tensor.py @@ -72,7 +72,11 @@ def astensor(array): # If the array is not allocated (empty) # we just create a new one if array.data.ptr == 0: - return torch.empty(array.shape, dtype=to_torch_dtype(array.dtype)) + return torch.empty( + array.shape, + dtype=to_torch_dtype(array.dtype), + device=array.device.id + ) return torch.as_tensor( _ArrayWithCudaArrayInterfaceHavingStrides(array), device=array.device.id, diff --git a/tests/test_tensor.py b/tests/test_tensor.py index fb33757..3a6c4be 100644 --- a/tests/test_tensor.py +++ b/tests/test_tensor.py @@ -59,22 +59,26 @@ def test_astensor_negative_stride(): def test_asarray_empty_cpu(): t = torch.tensor([], dtype=torch.float32) a = tensor.asarray(t) + assert isinstance(a, numpy.ndarray) def test_asarray_empty_gpu(): t = torch.tensor([], dtype=torch.float32, device='cuda') a = tensor.asarray(t) + assert isinstance(a, cupy.ndarray) def test_astensor_empty_cpu(): a = numpy.array([], dtype=numpy.float32) t = tensor.astensor(a) + assert t.device.type == 'cpu' def test_astensor_empty_gpu(): a = cupy.array([], dtype=cupy.float32) t = tensor.astensor(a) assert isinstance(t, torch.Tensor) + assert t.device.type == 'cuda' t += 1 numpy.testing.assert_array_equal(a.get(), t.cpu().numpy())