add assertion

Author: deependujha

src/lightning/pytorch/callbacks/batch_size_finder.py

@@ -131,6 +131,8 @@ def __init__(
         if mode not in self.SUPPORTED_MODES:
             raise ValueError(f"`mode` should be either of {self.SUPPORTED_MODES}")
 
+        assert 0.0 <= margin < 1.0, f"`margin` should be between 0 and 1. Found {margin=}"
+
         self.optimal_batch_size: Optional[int] = init_val
         self._mode = mode
         self._steps_per_trial = steps_per_trial

src/lightning/pytorch/tuner/batch_size_scaling.py

@@ -67,6 +67,8 @@ def _scale_batch_size(
             when running on CPU or when automatic OOM detection is not available.
 
     """
+    assert 0.0 <= margin < 1.0, f"`margin` should be between 0 and 1. Found {margin=}"
+
     if trainer.fast_dev_run:
         rank_zero_warn("Skipping batch size scaler since `fast_dev_run` is enabled.")
         return None
@@ -244,6 +246,8 @@ def _run_binsearch_scaling(
     Hereafter, the batch size is further refined using a binary search
 
     """
+    assert 0.0 <= margin < 1.0, f"`margin` should be between 0 and 1. Found {margin=}"
+
     low = 1
     high = None
     count = 0

src/lightning/pytorch/tuner/tuning.py

@@ -86,6 +86,7 @@ def scale_batch_size(
         """
         _check_tuner_configuration(train_dataloaders, val_dataloaders, dataloaders, method)
         _check_scale_batch_size_configuration(self._trainer)
+        assert 0.0 <= margin < 1.0, f"`margin` should be between 0 and 1. Found {margin=}"
 
         # local import to avoid circular import
         from lightning.pytorch.callbacks.batch_size_finder import BatchSizeFinder