From e97c28a02bb789918a66ee1978e418065c7c0414 Mon Sep 17 00:00:00 2001 From: jjenniferdai <89552168+jjenniferdai@users.noreply.github.com> Date: Fri, 3 Sep 2021 07:49:58 -0700 Subject: [PATCH] Typing `tuner.auto_gpu_select` (#9292) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: ananthsub Co-authored-by: Nicki Skafte --- pyproject.toml | 1 + pytorch_lightning/tuner/auto_gpu_select.py | 8 +++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3b5cad8874..61b620d759 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ module = [ "pytorch_lightning.trainer.evaluation_loop", "pytorch_lightning.trainer.connectors.logger_connector.fx_validator", "pytorch_lightning.trainer.connectors.logger_connector.logger_connector", + "pytorch_lightning.tuner.auto_gpu_select", "pytorch_lightning.utilities.apply_func", "pytorch_lightning.utilities.argparse", "pytorch_lightning.utilities.cli", diff --git a/pytorch_lightning/tuner/auto_gpu_select.py b/pytorch_lightning/tuner/auto_gpu_select.py index ccf969196f..5ecc51dbc3 100644 --- a/pytorch_lightning/tuner/auto_gpu_select.py +++ b/pytorch_lightning/tuner/auto_gpu_select.py @@ -11,12 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from typing import List + import torch from pytorch_lightning.utilities.exceptions import MisconfigurationException -def pick_multiple_gpus(nb): +def pick_multiple_gpus(nb: int) -> List[int]: """ Raises: MisconfigurationException: @@ -30,14 +32,14 @@ def pick_multiple_gpus(nb): nb = torch.cuda.device_count() if nb == -1 else nb - picked = [] + picked: List[int] = [] for _ in range(nb): picked.append(pick_single_gpu(exclude_gpus=picked)) return picked -def pick_single_gpu(exclude_gpus: list): +def pick_single_gpu(exclude_gpus: List[int]) -> int: """ Raises: RuntimeError: