From d1a56fcd9d68e3d3706a190c950a777d3104750f Mon Sep 17 00:00:00 2001 From: Jeff Yang Date: Tue, 5 Jan 2021 14:49:14 -0800 Subject: [PATCH] [docs] add docstring in torch.cuda.get_device_properties (#49792) Summary: Fixes https://github.com/pytorch/pytorch/issues/49737 Added docstring in `torch.cuda.get_device_properties` Added the `Returns` in `torch.cuda.get_device_name` Pull Request resolved: https://github.com/pytorch/pytorch/pull/49792 Reviewed By: mruberry Differential Revision: D25784046 Pulled By: ngimel fbshipit-source-id: f88da02147f92c889398957fcaf22961d3bb1062 --- torch/cuda/__init__.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py index 8ee83fa81fe7..5535cef78395 100644 --- a/torch/cuda/__init__.py +++ b/torch/cuda/__init__.py @@ -271,6 +271,9 @@ def get_device_name(device: Optional[_device_t] = None) -> str: name. This function is a no-op if this argument is a negative integer. It uses the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). + + Returns: + str: the name of the device """ return get_device_properties(device).name @@ -293,6 +296,15 @@ def get_device_capability(device: Optional[_device_t] = None) -> Tuple[int, int] def get_device_properties(device: _device_t) -> _CudaDeviceProperties: + r"""Gets the properties of a device. + + Args: + device (torch.device or int or str): device for which to return the + properties of the device. + + Returns: + _CudaDeviceProperties: the properties of the device + """ _lazy_init() # will define _get_device_properties device = _get_device_index(device, optional=True) if device < 0 or device >= device_count():