diff --git a/imperative/python/megengine/device.py b/imperative/python/megengine/device.py index 731b8985e692ba38a94366311a26dc4db3b08fb7..631e49d74cb05118c23102adafc37c2f42cade86 100644 --- a/imperative/python/megengine/device.py +++ b/imperative/python/megengine/device.py @@ -70,7 +70,7 @@ def set_default_device(device: str = "xpux"): multi-threading parallelism at the operator level. For example, 'multithread4' will compute with 4 threads. which implements - The default value is 'xpux' to specify any device available. + The default value is 'xpux' to specify any device available. The priority of using gpu is higher when both gpu and cpu are available. It can also be set by environmental variable `MGE_DEFAULT_DEVICE`. """ diff --git a/imperative/python/megengine/tensor.py b/imperative/python/megengine/tensor.py index 1848463c39605ee460ccdad7f555e178735a7648..89436323226ea8e10451b7a7ace36e90cef12277 100644 --- a/imperative/python/megengine/tensor.py +++ b/imperative/python/megengine/tensor.py @@ -11,6 +11,8 @@ import collections from .core import Tensor as _Tensor +from .core.ops.builtin import Copy +from .core.tensor.core import apply from .device import get_default_device @@ -30,6 +32,9 @@ class Tensor(_Tensor): def reset_zero(self): self *= 0 + def to(self, cn): + return apply(Copy(comp_node=cn), self)[0] + def __getstate__(self): r""" __getstate__ will be called for pickle serialization or deep copy """ diff --git a/imperative/python/test/unit/functional/test_tensor.py b/imperative/python/test/unit/functional/test_tensor.py index 72e1fb73bb0448690d3e35fbc088c4a27e170b12..2b86933d05fc1cbe6d9a004ce4c4611ac977a26e 100644 --- a/imperative/python/test/unit/functional/test_tensor.py +++ b/imperative/python/test/unit/functional/test_tensor.py @@ -322,6 +322,8 @@ def copy_test(dst, src): x = tensor(data, device=src) y = F.copy(x, dst) assert np.allclose(data, y.numpy()) + z = x.to(dst) + assert np.allclose(data, z.numpy()) @pytest.mark.skipif(