Move a Tensor to a Specific Device in PyTorch

Things on this page are fragmentary and immature notes/thoughts of the author. Please read with your own judgement!

import torch

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

tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

Tips¶

The methods Tensor.cpu, Tensor.cuda and Tensor.to are not in-palce. Instead, they return new copies of Tensors!
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¶

x.cpu()

tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

Tensor.cuda¶

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¶

x.to("cpu")

tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

x.to(torch.device("cpu:0"))

tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10]])

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

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