Fix batch_arg_name bug

Add `batch_arg_name` to all calls to `_adjust_batch_size`
Lightning-AI · Nov 22, 2020 · 2fb4f6c · 2fb4f6c
1 parent 8601268
commit 2fb4f6c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/tuner/batch_size_scaling.py b/pytorch_lightning/tuner/batch_size_scaling.py
@@ -102,7 +102,7 @@ def scale_batch_size(trainer,
         trainer.progress_bar_callback.disable()
 
     # Initially we just double in size until an OOM is encountered
-    new_size = _adjust_batch_size(trainer, value=init_val)  # initially set to init_val
+    new_size = _adjust_batch_size(trainer, batch_arg_name, value=init_val)  # initially set to init_val
     if mode == 'power':
         new_size = _run_power_scaling(trainer, model, new_size, batch_arg_name, max_trials, **fit_kwargs)
     elif mode == 'binsearch':
@@ -231,7 +231,7 @@ def _run_binsearch_scaling(trainer, model, new_size, batch_arg_name, max_trials,
                 garbage_collection_cuda()
                 high = new_size
                 midval = (high + low) // 2
-                new_size, _ = _adjust_batch_size(trainer, value=midval, desc='failed')
+                new_size, _ = _adjust_batch_size(trainer, batch_arg_name, value=midval, desc='failed')
                 if high - low <= 1:
                     break
             else: