Fix type hints

Author: ar90n

pyproject.toml

@@ -89,7 +89,6 @@ module = [
     "pytorch_lightning.trainer.optimizers",
     "pytorch_lightning.trainer.supporters",
     "pytorch_lightning.trainer.trainer",
-    "pytorch_lightning.tuner.batch_size_scaling",
     "pytorch_lightning.tuner.lr_finder",
     "pytorch_lightning.tuner.tuning",
     "pytorch_lightning.utilities.auto_restart",

src/pytorch_lightning/trainer/trainer.py

@@ -601,7 +601,7 @@ def _init_debugging_flags(
                 "Logging and checkpointing is suppressed."
             )
 
-        self.limit_train_batches = _determine_batch_limits(limit_train_batches, "limit_train_batches")
+        self.limit_train_batches: Union[int, float] = _determine_batch_limits(limit_train_batches, "limit_train_batches")
         self.limit_val_batches = _determine_batch_limits(limit_val_batches, "limit_val_batches")
         self.limit_test_batches = _determine_batch_limits(limit_test_batches, "limit_test_batches")
         self.limit_predict_batches = _determine_batch_limits(limit_predict_batches, "limit_predict_batches")

src/pytorch_lightning/tuner/batch_size_scaling.py

@@ -14,12 +14,14 @@
 import logging
 import os
 import uuid
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Union, List, Optional, Tuple
 
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
-from pytorch_lightning.loggers.logger import DummyLogger
+from pytorch_lightning.callbacks.callback import Callback
+from pytorch_lightning.loggers.logger import DummyLogger, Logger
+
 from pytorch_lightning.utilities.data import has_len_all_ranks
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.memory import garbage_collection_cuda, is_oom_error
@@ -41,7 +43,7 @@ def scale_batch_size(
     """See :meth:`~pytorch_lightning.tuner.tuning.Tuner.scale_batch_size`"""
     if trainer.fast_dev_run:
         rank_zero_warn("Skipping batch size scaler since fast_dev_run is enabled.")
-        return
+        return None
 
     if not lightning_hasattr(model, batch_arg_name):
         raise MisconfigurationException(f"Field {batch_arg_name} not found in both `model` and `model.hparams`")
@@ -234,18 +236,26 @@ def _adjust_batch_size(
     """
     model = trainer.lightning_module
     batch_size = lightning_getattr(model, batch_arg_name)
-    new_size = value if value is not None else int(batch_size * factor)
+    if value is not None:
+        new_size = value
+    else:
+        if not isinstance(batch_size, int):
+            raise ValueError(f"value is None and batch_size is not int value: {batch_size}")
+        new_size = int(batch_size * factor)
+
     if desc:
         log.info(f"Batch size {batch_size} {desc}, trying batch size {new_size}")
 
     if not _is_valid_batch_size(new_size, trainer.train_dataloader, trainer):
+        if not isinstance(trainer.train_dataloader, DataLoader):
+            raise ValueError("train_dataloader is not a DataLoader")
         new_size = min(new_size, len(trainer.train_dataloader.dataset))
 
     changed = new_size != batch_size
     lightning_setattr(model, batch_arg_name, new_size)
     return new_size, changed
 
 
-def _is_valid_batch_size(batch_size: int, dataloader: DataLoader, trainer: "pl.Trainer"):
+def _is_valid_batch_size(batch_size: int, dataloader: DataLoader, trainer: "pl.Trainer") -> bool:
     module = trainer.lightning_module or trainer.datamodule
     return not has_len_all_ranks(dataloader, trainer.strategy, module) or batch_size <= len(dataloader)