Ben Chuanlong Du's Blog

It is never too late to learn.

Move a Tensor to a Specific Device in PyTorch

In [1]:
import torch
In [2]:
x = torch.tensor(
    [
        [1, 2, 3, 4, 5],
        [6, 7, 8, 9, 10],
    ]
)
x
Out[2]:
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

Tips

  1. The methods Tensor.cpu, Tensor.cuda and Tensor.to are not in-palce. Instead, they return new copies of Tensors!

  2. There are basicially 2 ways to move a tensor and a module (notice that a model is a model too) to a specific device in PyTorch. The first (old) way is to call the methods Tensor.cpu and/or Tensor.cuda. The second (new) way is to call the method Tensor.to. Tensor.to is preferred over Tensor.cpu/Tensor.cuda as it is more flexible while almost as easy to use.

Tensor.cpu

In [5]:
x.cpu()
Out[5]:
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

Tensor.cuda

In [6]:
x.cuda()

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-6-795404c7d5aa> in <module>
----> 1 x.cuda()

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _lazy_init()
    194             raise RuntimeError(
    195                 "Cannot re-initialize CUDA in forked subprocess. " + msg)
--> 196         _check_driver()
    197         torch._C._cuda_init()
    198         _cudart = _load_cudart()

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _check_driver()
     92 def _check_driver():
     93     if not hasattr(torch._C, '_cuda_isDriverSufficient'):
---> 94         raise AssertionError("Torch not compiled with CUDA enabled")
     95     if not torch._C._cuda_isDriverSufficient():
     96         if torch._C._cuda_getDriverVersion() == 0:

AssertionError: Torch not compiled with CUDA enabled

Tensor.to

In [3]:
x.to("cpu")
Out[3]:
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
In [9]:
x.to(torch.device("cpu:0"))
Out[9]:
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])
In [4]:
x.to("cuda")

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-4-192d0d825d26> in <module>
----> 1 x.to("cuda")

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _lazy_init()
    194             raise RuntimeError(
    195                 "Cannot re-initialize CUDA in forked subprocess. " + msg)
--> 196         _check_driver()
    197         torch._C._cuda_init()
    198         _cudart = _load_cudart()

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _check_driver()
     92 def _check_driver():
     93     if not hasattr(torch._C, '_cuda_isDriverSufficient'):
---> 94         raise AssertionError("Torch not compiled with CUDA enabled")
     95     if not torch._C._cuda_isDriverSufficient():
     96         if torch._C._cuda_getDriverVersion() == 0:

AssertionError: Torch not compiled with CUDA enabled
In [10]:
x.to(torch.device("cuda:1"))

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-10-2ab8d6370486> in <module>
----> 1 x.to(torch.device("cuda:1"))

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _lazy_init()
    194             raise RuntimeError(
    195                 "Cannot re-initialize CUDA in forked subprocess. " + msg)
--> 196         _check_driver()
    197         torch._C._cuda_init()
    198         _cudart = _load_cudart()

~/.local/lib/python3.7/site-packages/torch/cuda/__init__.py in _check_driver()
     92 def _check_driver():
     93     if not hasattr(torch._C, '_cuda_isDriverSufficient'):
---> 94         raise AssertionError("Torch not compiled with CUDA enabled")
     95     if not torch._C._cuda_isDriverSufficient():
     96         if torch._C._cuda_getDriverVersion() == 0:

AssertionError: Torch not compiled with CUDA enabled

References

https://pytorch.org/docs/stable/tensors.html#torch.Tensor.device

In [ ]:

Comments