diff --git a/docs/source/tensor_attributes.rst b/docs/source/tensor_attributes.rst index 47becf990d5..c6592476d43 100644 --- a/docs/source/tensor_attributes.rst +++ b/docs/source/tensor_attributes.rst @@ -139,7 +139,7 @@ torch.device A :class:`torch.device` is an object representing the device on which a :class:`torch.Tensor` is or will be allocated. -The :class:`torch.device` contains a device type (``'cpu'`` or ``'cuda'``) and optional device +The :class:`torch.device` contains a device type (``'cpu'``, ``'cuda'`` or ``'mps'``) and optional device ordinal for the device type. If the device ordinal is not present, this object will always represent the current device for the device type, even after :func:`torch.cuda.set_device()` is called; e.g., a :class:`torch.Tensor` constructed with device ``'cuda'`` is equivalent to ``'cuda:X'`` where X is @@ -158,6 +158,9 @@ Via a string: >>> torch.device('cpu') device(type='cpu') + >>> torch.device('mps') + device(type='mps') + >>> torch.device('cuda') # current cuda device device(type='cuda') @@ -168,6 +171,9 @@ Via a string and device ordinal: >>> torch.device('cuda', 0) device(type='cuda', index=0) + >>> torch.device('mps', 0) + device(type='mps', index=0) + >>> torch.device('cpu', 0) device(type='cpu', index=0)