Repository: Lightning-AI/pytorch-lightning
Pull request: Fix accumulate_grad_batches for last batch
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
-->

Fixes #1549 

**EDIT:**

### Issue to reproduce – Colab
https://colab.research.google.com/drive/1iMMKPMeghhfWD092lT3l9XKClY2YJYNQ?usp=sharing

### Fixed issue to reproduce – Colab
https://colab.research.google.com/drive/1udmCrqYfJz2t8pEDi5pAIJKcyH_jHLLF?usp=sharing


# Before submitting

- [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
- [x] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests? 
- [x] Did you verify new and existing tests pass locally with your changes?
- [x] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Unwanted accumulate_grad_batches behavior
## 🐛 Bug

When using the flag `accumulate_grad_batches` for the trainer, if an action is to be performed at the last mini-batch it isn't done.

### To Reproduce

Steps to reproduce the behavior:

1. In `on_after_backward`, define some logic if we are at the last mini-batch
``` 
if self.__nbbatch -1 <= self.__batchidx:
 some_param.grad += gradient_penalty
```
2. Run with and without the flag
3. If running with the flag, the gradient penalty has not effect (the optimizer probably didn't take a step for the last mini-batch)

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
See above. 
### Expected behavior

I manually compute a gradient penalty that is applied only at the last mini-batch of an epoch. Using the flag shouldn't break this behavior.


<!-- A clear and concise description of what you expected to happen. -->

### Environment

PyTorch version: 1.6.0.dev20200403
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: CentOS Linux release 7.7.1908 (Core)
GCC version: (Homebrew GCC 5.5.0_7) 5.5.0
CMake version: version 3.13.0

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: Tesla P100-SXM2-16GB
GPU 1: Tesla P100-SXM2-16GB
GPU 2: Tesla P100-SXM2-16GB
GPU 3: Tesla P100-SXM2-16GB

Nvidia driver version: 440.33.01
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.18.1
[conda] blas                      1.0                         mkl  
[conda] kmeans-pytorch            0.2                      pypi_0    pypi
[conda] mkl                       2019.4                      243  
[conda] mkl-service               2.3.0            py37he904b0f_0  
[conda] mkl_fft                   1.0.15           py37ha843d7b_0  
[conda] mkl_random                1.1.0            py37hd6b4f25_0  
[conda] pytorch                   1.6.0.dev20200403 py3.7_cuda10.1.243_cudnn7.6.3_0    pytorch-nightly
[conda] pytorch-lightning         0.7.3                    pypi_0    pypi
[conda] pytorch-memlab            0.0.4                    pypi_0    pypi
[conda] pytorch-pcen              0.0.1                    pypi_0    pypi
[conda] torchvision               0.6.0.dev20200403      py37_cu101    pytorch-nightly

### Additional context

It isn't a bug per se, but it should at least be a documented behavior, ideally controllable with a flag.

Hints from issue discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

@sebastienwood thanks for bringing this up! we're looking at it for next sprint

Commit message: fixed with is_last_batch suggested by @awaelchli

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ad17f9f4ad482..793200363abab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -68,6 +68,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `accumulate_grad_batches` for last batch ([#2853](https://github.com/PyTorchLightning/pytorch-lightning/pull/2853))
+
 - Fixed setup call while testing ([#2624](https://github.com/PyTorchLightning/pytorch-lightning/pull/2624))
 
 - Fixed local rank zero casting ([#2640](https://github.com/PyTorchLightning/pytorch-lightning/pull/2640))
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
index 72178b8e8e17a..189b0ae0039c2 100644
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -368,9 +368,11 @@ def train(self):
                 if self.reload_dataloaders_every_epoch:
                     self.reset_train_dataloader(model)
                 # set seed for distributed sampler (enables shuffling for each epoch)
-                if (self.use_ddp or self.use_horovod or self.on_tpu) \
-                        and hasattr(self.train_dataloader, 'sampler') \
-                        and hasattr(self.train_dataloader.sampler, 'set_epoch'):
+                if (
+                    (self.use_ddp or self.use_horovod or self.on_tpu)
+                    and hasattr(self.train_dataloader, 'sampler')
+                    and hasattr(self.train_dataloader.sampler, 'set_epoch')
+                ):
                     self.train_dataloader.sampler.set_epoch(epoch)
 
                 # update training progress in trainer and model
@@ -381,9 +383,7 @@ def train(self):
                 self.accumulation_scheduler.on_epoch_start(self, self.get_model())
 
                 # stores accumulated grad fractions per batch
-                self.batch_loss_value = TensorRunningAccum(
-                    window_length=self.accumulate_grad_batches
-                )
+                self.batch_loss_value = TensorRunningAccum(window_length=self.accumulate_grad_batches)
 
                 # -----------------
                 # RUN TNG EPOCH
@@ -402,13 +402,15 @@ def train(self):
                 met_min_steps = self.global_step >= self.min_steps if self.min_steps else True
 
                 if self.should_stop:
-                    if (met_min_epochs and met_min_steps):
+                    if met_min_epochs and met_min_steps:
                         self.run_training_teardown()
                         return
                     else:
-                        log.info('Trainer was signaled to stop but required minimum epochs'
-                                 f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has'
-                                 ' not been met. Training will continue...')
+                        log.info(
+                            'Trainer was signaled to stop but required minimum epochs'
+                            f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has'
+                            ' not been met. Training will continue...'
+                        )
 
             self.run_training_teardown()
 
@@ -471,13 +473,14 @@ def run_training_epoch(self):
 
         # run epoch
         for batch_idx, (batch, is_last_batch) in self.profiler.profile_iterable(
-                enumerate(_with_is_last(train_dataloader)), "get_train_batch"
+            enumerate(_with_is_last(train_dataloader)), "get_train_batch"
         ):
             # stop epoch if we limited the number of training batches
             if batch_idx >= self.num_training_batches:
                 break
 
             self.batch_idx = batch_idx
+            self.is_last_batch = is_last_batch
             model.global_step = self.global_step
 
             # ------------------------------------
@@ -555,7 +558,7 @@ def check_checkpoint_callback(self, should_check_val):
             [c.on_validation_end(self, self.get_model()) for c in checkpoint_callbacks]
 
     def update_train_loop_lr_schedulers(self):
-        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0 or self.is_last_batch:
             # update lr
             self.update_learning_rates(interval='step')
 
@@ -643,7 +646,7 @@ def sync_horovod(self):
 
     def increment_accumulated_grad_global_step(self):
         # progress global step according to grads progress
-        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+        if (self.batch_idx + 1) % self.accumulate_grad_batches == 0 or self.is_last_batch:
             self.global_step += 1
         self.total_batch_idx += 1
 
@@ -670,7 +673,7 @@ def should_check_val(self, batch_idx, is_last_batch):
         can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
         can_check_val = self.enable_validation and can_check_epoch
         should_check_val = is_val_check_batch or self.should_stop
-        is_last_batch_for_infinite_dataset = (is_last_batch and self.val_check_batch == float('inf'))
+        is_last_batch_for_infinite_dataset = is_last_batch and self.val_check_batch == float('inf')
         should_check_val = can_check_val and (should_check_val or is_last_batch_for_infinite_dataset)
 
         return should_check_val
@@ -778,7 +781,7 @@ def run_training_batch(self, batch, batch_idx):
                 # BACKWARD PASS
                 # ------------------------------
                 # gradient update with accumulated gradients
-                if (self.batch_idx + 1) % self.accumulate_grad_batches == 0:
+                if (self.batch_idx + 1) % self.accumulate_grad_batches == 0 or self.is_last_batch:
 
                     # backward
                     grad_norm_dic = self.run_batch_backward_pass(split_batch, batch_idx, opt_idx, optimizer)
@@ -816,7 +819,7 @@ def run_training_batch(self, batch, batch_idx):
             signal=0,
             grad_norm_dic=grad_norm_dic,
             batch_log_metrics=batch_log_metrics,
-            training_step_output_for_epoch_end=opt_closure_result.training_step_output_for_epoch_end
+            training_step_output_for_epoch_end=opt_closure_result.training_step_output_for_epoch_end,
         )
         return result
 
@@ -829,8 +832,7 @@ def run_batch_backward_pass(self, split_batch, batch_idx, opt_idx, optimizer):
         if batch_idx % self.row_log_interval == 0:
             if float(self.track_grad_norm) > 0:
                 model = self.get_model()
-                grad_norm_dic = model.grad_norm(
-                    self.track_grad_norm)
+                grad_norm_dic = model.grad_norm(self.track_grad_norm)
 
         # ------------------
         # CLIP GRADS
@@ -862,8 +864,7 @@ def call_optimizer_step(self, optimizer, opt_idx, batch_idx, split_batch):
 
             # apply TPU optimizer
             if self.use_tpu and XLA_AVAILABLE:
-                model.optimizer_step(self.current_epoch, batch_idx,
-                                     optimizer, opt_idx, lambda_closure, on_tpu=True)
+                model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, on_tpu=True)
 
             # for LBFGS do something a bit different
             elif isinstance(optimizer, torch.optim.LBFGS):
@@ -872,9 +873,11 @@ def call_optimizer_step(self, optimizer, opt_idx, batch_idx, split_batch):
                 if self.amp_type == AMPType.NATIVE:
                     raise MisconfigurationException(
                         'native PyTorch amp and lbfgs are not compatible.'
-                        ' To request, please file a Github issue in PyTorch and tag @mcarilli')
-                model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure,
-                                     using_lbfgs=True)
+                        ' To request, please file a Github issue in PyTorch and tag @mcarilli'
+                    )
+                model.optimizer_step(
+                    self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, using_lbfgs=True
+                )
 
             # when using 16-bit
             else:
@@ -902,11 +905,9 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
         with self.profiler.profile('model_forward'):
             if self.amp_type == AMPType.NATIVE and not self.use_tpu:
                 with torch.cuda.amp.autocast():
-                    training_step_output = self.training_forward(split_batch, batch_idx,
-                                                                 opt_idx, hiddens)
+                    training_step_output = self.training_forward(split_batch, batch_idx, opt_idx, hiddens)
             else:
-                training_step_output = self.training_forward(split_batch, batch_idx, opt_idx,
-                                                             hiddens)
+                training_step_output = self.training_forward(split_batch, batch_idx, opt_idx, hiddens)
 
             # ----------------------------
             # PROCESS THE RESULT
@@ -917,8 +918,9 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
 
             # don't allow EvalResult in the training_step
             if isinstance(training_step_output, EvalResult):
-                raise MisconfigurationException('training_step cannot return EvalResult, '
-                                                'use a dict or TrainResult instead')
+                raise MisconfigurationException(
+                    'training_step cannot return EvalResult, ' 'use a dict or TrainResult instead'
+                )
 
             # handle regular dicts
             if not is_result_obj:
@@ -1125,8 +1127,10 @@ def training_forward(self, batch, batch_idx, opt_idx, hiddens):
             with self.profiler.profile('training_end'):
                 output = model_ref.training_end(output)
 
-            rank_zero_warn('`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.'
-                           ' Use training_epoch_end instead', DeprecationWarning)
+            rank_zero_warn(
+                '`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.' ' Use training_epoch_end instead',
+                DeprecationWarning,
+            )
 
         return output
 
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 2444d9905f6a1..a16ee1801765e 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -16,8 +16,7 @@
 import tests.base.develop_utils as tutils
 from pytorch_lightning import Callback, LightningModule, Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-from pytorch_lightning.core.saving import (
-    load_hparams_from_tags_csv, load_hparams_from_yaml, save_hparams_to_tags_csv)
+from pytorch_lightning.core.saving import load_hparams_from_tags_csv, load_hparams_from_yaml, save_hparams_to_tags_csv
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.utilities.cloud_io import load as pl_load
@@ -37,10 +36,7 @@ def test_no_val_module(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
     logger = tutils.get_default_logger(tmpdir)
 
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir),
+        default_root_dir=tmpdir, max_epochs=1, logger=logger, checkpoint_callback=ModelCheckpoint(tmpdir),
     )
     # fit model
     result = trainer.fit(model)
@@ -58,11 +54,12 @@ def test_no_val_module(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
     # load new model
     hparams_path = tutils.get_data_path(logger, path_dir=tmpdir)
     hparams_path = os.path.join(hparams_path, 'hparams.yaml')
-    ckpt_path = f'http://{tmpdir_server[0]}:{tmpdir_server[1]}/{os.path.basename(new_weights_path)}' if url_ckpt else new_weights_path
-    model_2 = EvalModelTemplate.load_from_checkpoint(
-        checkpoint_path=ckpt_path,
-        hparams_file=hparams_path,
+    ckpt_path = (
+        f'http://{tmpdir_server[0]}:{tmpdir_server[1]}/{os.path.basename(new_weights_path)}'
+        if url_ckpt
+        else new_weights_path
     )
+    model_2 = EvalModelTemplate.load_from_checkpoint(checkpoint_path=ckpt_path, hparams_file=hparams_path,)
     model_2.eval()
 
 
@@ -79,10 +76,7 @@ def test_no_val_end_module(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
 
     # fit model
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        logger=logger,
-        checkpoint_callback=ModelCheckpoint(tmpdir),
+        default_root_dir=tmpdir, max_epochs=1, logger=logger, checkpoint_callback=ModelCheckpoint(tmpdir),
     )
     result = trainer.fit(model)
 
@@ -96,21 +90,18 @@ def test_no_val_end_module(monkeypatch, tmpdir, tmpdir_server, url_ckpt):
     # load new model
     hparams_path = tutils.get_data_path(logger, path_dir=tmpdir)
     hparams_path = os.path.join(hparams_path, 'hparams.yaml')
-    ckpt_path = f'http://{tmpdir_server[0]}:{tmpdir_server[1]}/{os.path.basename(new_weights_path)}' if url_ckpt else new_weights_path
-    model_2 = EvalModelTemplate.load_from_checkpoint(
-        checkpoint_path=ckpt_path,
-        hparams_file=hparams_path,
+    ckpt_path = (
+        f'http://{tmpdir_server[0]}:{tmpdir_server[1]}/{os.path.basename(new_weights_path)}'
+        if url_ckpt
+        else new_weights_path
     )
+    model_2 = EvalModelTemplate.load_from_checkpoint(checkpoint_path=ckpt_path, hparams_file=hparams_path,)
     model_2.eval()
 
 
 @pytest.mark.parametrize(
     ['schedule', 'expected'],
-    [
-        pytest.param({1: 2, 3: 4}, [1, 2, 4]),
-        pytest.param(3, [3, 3, 3]),
-        pytest.param(4, [4, 4, 4])
-    ]
+    [pytest.param({1: 2, 3: 4}, [1, 2, 4]), pytest.param(3, [3, 3, 3]), pytest.param(4, [4, 4, 4])],
 )
 def test_gradient_accumulation_scheduling(tmpdir, schedule, expected):
     """
@@ -128,7 +119,7 @@ def test_gradient_accumulation_scheduling(tmpdir, schedule, expected):
     with pytest.raises(TypeError):
         assert Trainer(accumulate_grad_batches=[[2, 3], [4, 6]])
     with pytest.raises(TypeError):
-        assert Trainer(accumulate_grad_batches={1: 2, 3.: 4})
+        assert Trainer(accumulate_grad_batches={1: 2, 3.0: 4})
     with pytest.raises(TypeError):
         assert Trainer(accumulate_grad_batches={1: 2.5, 3: 5})
 
@@ -143,9 +134,16 @@ def test_gradient_accumulation_scheduling(tmpdir, schedule, expected):
     )
 
     # test optimizer call freq matches scheduler
-    def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
-                        second_order_closure=None, on_tpu=False,
-                        using_native_amp=False, using_lbfgs=False):
+    def _optimizer_step(
+        epoch,
+        batch_idx,
+        optimizer,
+        optimizer_idx,
+        second_order_closure=None,
+        on_tpu=False,
+        using_native_amp=False,
+        using_lbfgs=False,
+    ):
         # only test the first 12 batches in epoch
         if batch_idx < 12:
             if epoch == 0:
@@ -192,6 +190,55 @@ def _optimizer_step(epoch, batch_idx, optimizer, optimizer_idx,
     trainer.fit(model)
 
 
+@pytest.mark.parametrize(
+    ['schedule', 'limit'],
+    [
+        pytest.param({1: 20, 3: 40}, 1.0),
+        pytest.param({1: 20, 3: 40}, 0.5),
+        pytest.param(30, 1.0),
+        pytest.param(30, 0.6),
+        pytest.param(40, 1.0),
+        pytest.param(40, 0.7),
+    ],
+)
+def test_gradient_accumulation_scheduling_last_batch(tmpdir, schedule, limit):
+    """
+    Verify optimizer.step() applied to last batch while grad accumulation
+    """
+
+    model = EvalModelTemplate()
+
+    trainer = Trainer(
+        accumulate_grad_batches=schedule, max_epochs=4, limit_train_batches=limit, default_root_dir=tmpdir,
+    )
+
+    loss_backward = [0] * len(list(model.parameters()))
+    opt_step = [torch.tensor([0.0])] * len(list(model.parameters()))
+
+    def on_after_backward_():
+        for i, param in enumerate(model.parameters()):
+            loss_backward[i] = param.clone().data
+
+    def on_before_zero_grad_(optimizer):
+        for i, param in enumerate(model.parameters()):
+            opt_step[i] = param.clone().data
+
+    def on_train_batch_end_(batch, batch_idx, dataloader_idx):
+        # to check optimizer.step() get called or not
+        for i, (loss_param, opt_param) in enumerate(zip(loss_backward, opt_step)):
+            if trainer.current_epoch == 0:  # for accumulate_grad_batches={0: 1}
+                assert torch.equal(loss_param, opt_param) is False
+            if trainer.is_last_batch:
+                # model.parameters() get updated
+                assert torch.equal(loss_param, opt_param) is False
+
+    # for the test
+    model.on_after_backward = on_after_backward_
+    model.on_before_zero_grad = on_before_zero_grad_
+    model.on_train_batch_end = on_train_batch_end_
+    trainer.fit(model)
+
+
 def test_loading_meta_tags(tmpdir):
     """ test for backward compatibility to meta_tags.csv """
     tutils.reset_seed()
@@ -248,31 +295,37 @@ def test_dp_output_reduce():
     assert mixin.reduce_distributed_output(out, num_gpus=2) == out.mean()
 
     # when we have a dict of vals
-    out = {
-        'a': out,
-        'b': {
-            'c': out
-        }
-    }
+    out = {'a': out, 'b': {'c': out}}
     reduced = mixin.reduce_distributed_output(out, num_gpus=3)
     assert reduced['a'] == out['a']
     assert reduced['b']['c'] == out['b']['c']
 
 
-@pytest.mark.parametrize(["save_top_k", "save_last", "file_prefix", "expected_files"], [
-    pytest.param(-1, False, '', {'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt', 'epoch=1.ckpt', 'epoch=0.ckpt'},
-                 id="CASE K=-1  (all)"),
-    pytest.param(1, False, 'test_prefix_', {'test_prefix_epoch=4.ckpt'},
-                 id="CASE K=1 (2.5, epoch 4)"),
-    pytest.param(2, False, '', {'epoch=4.ckpt', 'epoch=2.ckpt'},
-                 id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"),
-    pytest.param(4, False, '', {'epoch=1.ckpt', 'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt'},
-                 id="CASE K=4 (save all 4 base)"),
-    pytest.param(3, False, '', {'epoch=2.ckpt', 'epoch=3.ckpt', 'epoch=4.ckpt'},
-                 id="CASE K=3 (save the 2nd, 3rd, 4th model)"),
-    pytest.param(1, True, '', {'epoch=4.ckpt', 'last.ckpt'},
-                 id="CASE K=1 (save the 4th model and the last model)"),
-])
+@pytest.mark.parametrize(
+    ["save_top_k", "save_last", "file_prefix", "expected_files"],
+    [
+        pytest.param(
+            -1,
+            False,
+            '',
+            {'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt', 'epoch=1.ckpt', 'epoch=0.ckpt'},
+            id="CASE K=-1  (all)",
+        ),
+        pytest.param(1, False, 'test_prefix_', {'test_prefix_epoch=4.ckpt'}, id="CASE K=1 (2.5, epoch 4)"),
+        pytest.param(2, False, '', {'epoch=4.ckpt', 'epoch=2.ckpt'}, id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"),
+        pytest.param(
+            4,
+            False,
+            '',
+            {'epoch=1.ckpt', 'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt'},
+            id="CASE K=4 (save all 4 base)",
+        ),
+        pytest.param(
+            3, False, '', {'epoch=2.ckpt', 'epoch=3.ckpt', 'epoch=4.ckpt'}, id="CASE K=3 (save the 2nd, 3rd, 4th model)"
+        ),
+        pytest.param(1, True, '', {'epoch=4.ckpt', 'last.ckpt'}, id="CASE K=1 (save the 4th model and the last model)"),
+    ],
+)
 def test_model_checkpoint_options(tmpdir, save_top_k, save_last, file_prefix, expected_files):
     """Test ModelCheckpoint options."""
 
@@ -282,8 +335,9 @@ def mock_save_function(filepath, *args):
     # simulated losses
     losses = [10, 9, 2.8, 5, 2.5]
 
-    checkpoint_callback = ModelCheckpoint(tmpdir, save_top_k=save_top_k, save_last=save_last,
-                                          prefix=file_prefix, verbose=1)
+    checkpoint_callback = ModelCheckpoint(
+        tmpdir, save_top_k=save_top_k, save_last=save_last, prefix=file_prefix, verbose=1
+    )
     checkpoint_callback.save_function = mock_save_function
     trainer = Trainer()
 
@@ -295,8 +349,10 @@ def mock_save_function(filepath, *args):
 
     file_lists = set(os.listdir(tmpdir))
 
-    assert len(file_lists) == len(expected_files), \
-        "Should save %i models when save_top_k=%i" % (len(expected_files), save_top_k)
+    assert len(file_lists) == len(expected_files), "Should save %i models when save_top_k=%i" % (
+        len(expected_files),
+        save_top_k,
+    )
 
     # verify correct naming
     for fname in expected_files:
@@ -310,9 +366,7 @@ def test_model_checkpoint_only_weights(tmpdir):
     model = EvalModelTemplate()
 
     trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        checkpoint_callback=ModelCheckpoint(tmpdir, save_weights_only=True),
+        default_root_dir=tmpdir, max_epochs=1, checkpoint_callback=ModelCheckpoint(tmpdir, save_weights_only=True),
     )
     # fit model
     result = trainer.fit(model)
@@ -391,7 +445,7 @@ def increment_on_load_checkpoint(self, _):
         checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1),
         default_root_dir=tmpdir,
         early_stop_callback=False,
-        val_check_interval=1.,
+        val_check_interval=1.0,
     )
 
     trainer = Trainer(**trainer_options)
@@ -432,9 +486,7 @@ def _init_steps_model():
     # get number of samples in 1 epoch
     num_train_samples = math.floor(len(model.train_dataloader()) * train_percent)
 
-    trainer_options = dict(
-        limit_train_batches=train_percent,
-    )
+    trainer_options = dict(limit_train_batches=train_percent,)
     return model, trainer_options, num_train_samples
 
 
@@ -444,9 +496,7 @@ def test_trainer_max_steps_and_epochs(tmpdir):
 
     # define less train steps than epochs
     trainer_options.update(
-        default_root_dir=tmpdir,
-        max_epochs=3,
-        max_steps=num_train_samples + 10,
+        default_root_dir=tmpdir, max_epochs=3, max_steps=num_train_samples + 10,
     )
 
     # fit model
@@ -459,8 +509,7 @@ def test_trainer_max_steps_and_epochs(tmpdir):
 
     # define less train epochs than steps
     trainer_options.update(
-        max_epochs=2,
-        max_steps=trainer_options['max_epochs'] * 2 * num_train_samples,
+        max_epochs=2, max_steps=trainer_options['max_epochs'] * 2 * num_train_samples,
     )
 
     # fit model
@@ -495,8 +544,9 @@ def test_trainer_min_steps_and_epochs(tmpdir):
     assert result == 1, "Training did not complete"
 
     # check model ran for at least min_epochs
-    assert trainer.global_step >= num_train_samples and \
-        trainer.current_epoch > 0, "Model did not train for at least min_epochs"
+    assert (
+        trainer.global_step >= num_train_samples and trainer.current_epoch > 0
+    ), "Model did not train for at least min_epochs"
 
     # define less epochs than min_steps
     trainer_options['min_steps'] = math.floor(num_train_samples * 1.5)
@@ -507,8 +557,9 @@ def test_trainer_min_steps_and_epochs(tmpdir):
     assert result == 1, "Training did not complete"
 
     # check model ran for at least num_train_samples*1.5
-    assert trainer.global_step >= math.floor(num_train_samples * 1.5) and \
-        trainer.current_epoch > 0, "Model did not train for at least min_steps"
+    assert (
+        trainer.global_step >= math.floor(num_train_samples * 1.5) and trainer.current_epoch > 0
+    ), "Model did not train for at least min_steps"
 
 
 def test_benchmark_option(tmpdir):
@@ -521,11 +572,7 @@ def test_benchmark_option(tmpdir):
     assert not torch.backends.cudnn.benchmark
 
     # fit model
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        benchmark=True,
-    )
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, benchmark=True,)
     result = trainer.fit(model)
 
     # verify training completed
@@ -567,7 +614,11 @@ def test_test_checkpoint_path(tmpdir, ckpt_path, save_top_k):
             with pytest.raises(FileNotFoundError):
                 trainer.test(ckpt_path='random.ckpt')
         else:
-            ckpt_path = str(list((Path(tmpdir) / f'lightning_logs/version_{trainer.logger.version}/checkpoints').iterdir())[0].absolute())
+            ckpt_path = str(
+                list((Path(tmpdir) / f'lightning_logs/version_{trainer.logger.version}/checkpoints').iterdir())[
+                    0
+                ].absolute()
+            )
             trainer.test(ckpt_path=ckpt_path)
             assert trainer.tested_ckpt_path == ckpt_path
 
@@ -606,10 +657,8 @@ def validation_epoch_end(self, *args, **kwargs):
     # check that limit_val_batches=0 turns off validation
     assert result == 1, 'training failed to complete'
     assert trainer.current_epoch == 1
-    assert not model.validation_step_invoked, \
-        '`validation_step` should not run when `limit_val_batches=0`'
-    assert not model.validation_epoch_end_invoked, \
-        '`validation_epoch_end` should not run when `limit_val_batches=0`'
+    assert not model.validation_step_invoked, '`validation_step` should not run when `limit_val_batches=0`'
+    assert not model.validation_epoch_end_invoked, '`validation_epoch_end` should not run when `limit_val_batches=0`'
 
     # check that limit_val_batches has no influence when fast_dev_run is turned on
     model = CurrentModel(**hparams)
@@ -619,14 +668,11 @@ def validation_epoch_end(self, *args, **kwargs):
 
     assert result == 1, 'training failed to complete'
     assert trainer.current_epoch == 0
-    assert model.validation_step_invoked, \
-        'did not run `validation_step` with `fast_dev_run=True`'
-    assert model.validation_epoch_end_invoked, \
-        'did not run `validation_epoch_end` with `fast_dev_run=True`'
+    assert model.validation_step_invoked, 'did not run `validation_step` with `fast_dev_run=True`'
+    assert model.validation_epoch_end_invoked, 'did not run `validation_epoch_end` with `fast_dev_run=True`'
 
 
 def test_nan_loss_detection(tmpdir):
-
     class CurrentModel(EvalModelTemplate):
         test_batch_inf_loss = 8
 
@@ -642,11 +688,7 @@ def training_step(self, batch, batch_idx, optimizer_idx=None):
     model = CurrentModel()
 
     # fit model
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_steps=(model.test_batch_inf_loss + 1),
-        terminate_on_nan=True,
-    )
+    trainer = Trainer(default_root_dir=tmpdir, max_steps=(model.test_batch_inf_loss + 1), terminate_on_nan=True,)
 
     with pytest.raises(ValueError, match=r'.*The loss returned in `training_step` is nan or inf.*'):
         trainer.fit(model)
@@ -657,7 +699,6 @@ def training_step(self, batch, batch_idx, optimizer_idx=None):
 
 
 def test_nan_params_detection(tmpdir):
-
     class CurrentModel(EvalModelTemplate):
         test_batch_nan = 8
 
@@ -667,11 +708,7 @@ def on_after_backward(self):
                 torch.nn.init.constant_(self.c_d1.bias, math.nan)
 
     model = CurrentModel()
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_steps=(model.test_batch_nan + 1),
-        terminate_on_nan=True,
-    )
+    trainer = Trainer(default_root_dir=tmpdir, max_steps=(model.test_batch_nan + 1), terminate_on_nan=True,)
 
     with pytest.raises(ValueError, match=r'.*Detected nan and/or inf values in `c_d1.bias`.*'):
         trainer.fit(model)
@@ -734,12 +771,7 @@ def _optimizer_step(*args, **kwargs):
         grad_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2)
         assert (grad_norm - 1.0).abs() < 0.01, "Gradient norm != 1.0: {grad_norm}".format(grad_norm=grad_norm)
 
-    trainer = Trainer(
-        max_steps=1,
-        max_epochs=1,
-        gradient_clip_val=1.0,
-        default_root_dir=tmpdir,
-    )
+    trainer = Trainer(max_steps=1, max_epochs=1, gradient_clip_val=1.0, default_root_dir=tmpdir,)
 
     # for the test
     model.optimizer_step = _optimizer_step
@@ -749,9 +781,7 @@ def _optimizer_step(*args, **kwargs):
 
 
 def test_gpu_choice(tmpdir):
-    trainer_options = dict(
-        default_root_dir=tmpdir,
-    )
+    trainer_options = dict(default_root_dir=tmpdir,)
     # Only run if CUDA is available
     if not torch.cuda.is_available():
         return
@@ -763,19 +793,22 @@ def test_gpu_choice(tmpdir):
         Trainer(**trainer_options, gpus=num_gpus + 1, auto_select_gpus=True)
 
 
-@pytest.mark.parametrize(['tpu_cores', 'expected_tpu_id', 'error_expected'], [
-    pytest.param(1, None, False),
-    pytest.param(8, None, False),
-    pytest.param([1], 1, False),
-    pytest.param([8], 8, False),
-    pytest.param('1,', 1, False),
-    pytest.param('1', None, False),
-    pytest.param('9, ', 9, True),
-    pytest.param([9], 9, True),
-    pytest.param([0], 0, True),
-    pytest.param(2, None, True),
-    pytest.param(10, None, True),
-])
+@pytest.mark.parametrize(
+    ['tpu_cores', 'expected_tpu_id', 'error_expected'],
+    [
+        pytest.param(1, None, False),
+        pytest.param(8, None, False),
+        pytest.param([1], 1, False),
+        pytest.param([8], 8, False),
+        pytest.param('1,', 1, False),
+        pytest.param('1', None, False),
+        pytest.param('9, ', 9, True),
+        pytest.param([9], 9, True),
+        pytest.param([0], 0, True),
+        pytest.param(2, None, True),
+        pytest.param(10, None, True),
+    ],
+)
 def test_tpu_choice(tmpdir, tpu_cores, expected_tpu_id, error_expected):
     if error_expected:
         with pytest.raises(MisconfigurationException, match=r'.*tpu_cores` can only be 1, 8 or [<1-8>]*'):
@@ -785,12 +818,10 @@ def test_tpu_choice(tmpdir, tpu_cores, expected_tpu_id, error_expected):
         assert trainer.tpu_id == expected_tpu_id
 
 
-@pytest.mark.parametrize(['limit_val_batches'], [
-    pytest.param(0.0),  # this should run no sanity checks
-    pytest.param(1),
-    pytest.param(1.0),
-    pytest.param(0.3),
-])
+@pytest.mark.parametrize(
+    ['limit_val_batches'],
+    [pytest.param(0.0), pytest.param(1), pytest.param(1.0), pytest.param(0.3),],  # this should run no sanity checks
+)
 def test_num_sanity_val_steps(tmpdir, limit_val_batches):
     """
     Test that num_sanity_val_steps=-1 runs through all validation data once.
@@ -813,85 +844,224 @@ def test_num_sanity_val_steps(tmpdir, limit_val_batches):
         assert mocked.call_count == sum(len(dl) * (limit_val_batches > 0) for dl in val_dataloaders)
 
 
-@pytest.mark.parametrize("trainer_kwargs,expected", [
-    pytest.param(
-        dict(distributed_backend=None, gpus=None),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend="dp", gpus=None),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend="dp", gpus=None),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp", gpus=None),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp", num_processes=2, gpus=None),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=2)
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp", num_nodes=2, gpus=None),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp_cpu", num_processes=2, gpus=None),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=2)
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp2", gpus=None),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=1)
-    ),
-    pytest.param(
-        dict(distributed_backend=None, gpus=1),
-        dict(use_dp=False, use_ddp=False, use_ddp2=False, num_gpus=1, on_gpu=True, use_single_gpu=True, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="dp", gpus=1),
-        dict(use_dp=True, use_ddp=False, use_ddp2=False, num_gpus=1, on_gpu=True, use_single_gpu=True, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp", gpus=1),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=1, on_gpu=True, use_single_gpu=True, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp_cpu", num_processes=2, gpus=1),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=0, on_gpu=False, use_single_gpu=False, num_processes=2),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp2", gpus=1),
-        dict(use_dp=False, use_ddp=False, use_ddp2=True, num_gpus=1, on_gpu=True, use_single_gpu=False, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend=None, gpus=2),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=2, on_gpu=True, use_single_gpu=False, num_processes=2),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="dp", gpus=2),
-        dict(use_dp=True, use_ddp=False, use_ddp2=False, num_gpus=2, on_gpu=True, use_single_gpu=False, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp", gpus=2),
-        dict(use_dp=False, use_ddp=True, use_ddp2=False, num_gpus=2, on_gpu=True, use_single_gpu=False, num_processes=2),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")]
-    ),
-    pytest.param(
-        dict(distributed_backend="ddp2", gpus=2),
-        dict(use_dp=False, use_ddp=False, use_ddp2=True, num_gpus=2, on_gpu=True, use_single_gpu=False, num_processes=1),
-        marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")]
-    ),
-])
+@pytest.mark.parametrize(
+    "trainer_kwargs,expected",
+    [
+        pytest.param(
+            dict(distributed_backend=None, gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="dp", gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="dp", gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp", gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp", num_processes=2, gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=2,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp", num_nodes=2, gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp_cpu", num_processes=2, gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=2,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp2", gpus=None),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+        ),
+        pytest.param(
+            dict(distributed_backend=None, gpus=1),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=1,
+                on_gpu=True,
+                use_single_gpu=True,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="dp", gpus=1),
+            dict(
+                use_dp=True,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=1,
+                on_gpu=True,
+                use_single_gpu=True,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp", gpus=1),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=1,
+                on_gpu=True,
+                use_single_gpu=True,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp_cpu", num_processes=2, gpus=1),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=0,
+                on_gpu=False,
+                use_single_gpu=False,
+                num_processes=2,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp2", gpus=1),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=True,
+                num_gpus=1,
+                on_gpu=True,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() == 0, reason="GPU needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend=None, gpus=2),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=2,
+                on_gpu=True,
+                use_single_gpu=False,
+                num_processes=2,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="dp", gpus=2),
+            dict(
+                use_dp=True,
+                use_ddp=False,
+                use_ddp2=False,
+                num_gpus=2,
+                on_gpu=True,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp", gpus=2),
+            dict(
+                use_dp=False,
+                use_ddp=True,
+                use_ddp2=False,
+                num_gpus=2,
+                on_gpu=True,
+                use_single_gpu=False,
+                num_processes=2,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")],
+        ),
+        pytest.param(
+            dict(distributed_backend="ddp2", gpus=2),
+            dict(
+                use_dp=False,
+                use_ddp=False,
+                use_ddp2=True,
+                num_gpus=2,
+                on_gpu=True,
+                use_single_gpu=False,
+                num_processes=1,
+            ),
+            marks=[pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Multiple GPUs needed")],
+        ),
+    ],
+)
 def test_trainer_config(trainer_kwargs, expected):
     trainer = Trainer(**trainer_kwargs)
     assert trainer.use_dp is expected["use_dp"]
@@ -908,7 +1078,6 @@ def test_trainer_subclassing():
 
     # First way of pulling out args from signature is to list them
     class TrainerSubclass(Trainer):
-
         def __init__(self, custom_arg, *args, custom_kwarg='test', **kwargs):
             super().__init__(*args, **kwargs)
             self.custom_arg = custom_arg
@@ -924,7 +1093,6 @@ def __init__(self, custom_arg, *args, custom_kwarg='test', **kwargs):
     # Second way is to pop from the dict
     # It's a special case because Trainer does not have any positional args
     class TrainerSubclass(Trainer):
-
         def __init__(self, **kwargs):
             self.custom_arg = kwargs.pop('custom_arg', 0)
             self.custom_kwarg = kwargs.pop('custom_kwarg', 'test')
@@ -941,20 +1109,17 @@ def __init__(self, **kwargs):
         TrainerSubclass(abcdefg='unknown_arg')
 
 
-@pytest.mark.parametrize('trainer_params', [
-    OmegaConf.create({'max_epochs': 1, 'gpus': 1}),
-    OmegaConf.create({'max_epochs': 1, 'gpus': [0]}),
-])
+@pytest.mark.parametrize(
+    'trainer_params',
+    [OmegaConf.create({'max_epochs': 1, 'gpus': 1}), OmegaConf.create({'max_epochs': 1, 'gpus': [0]}),],
+)
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
 def test_trainer_omegaconf(trainer_params):
     Trainer(**trainer_params)
 
 
 def test_trainer_pickle(tmpdir):
-    trainer = Trainer(
-        max_epochs=1,
-        default_root_dir=tmpdir,
-    )
+    trainer = Trainer(max_epochs=1, default_root_dir=tmpdir,)
     pickle.dumps(trainer)
     cloudpickle.dumps(trainer)
 
@@ -963,23 +1128,17 @@ def test_trainer_setup_call(tmpdir):
     """Test setup call with fit and test call."""
 
     class CurrentModel(EvalModelTemplate):
-
         def setup(self, stage):
             self.stage = stage
 
     class TrainerSubclass(Trainer):
-
         def setup(self, stage):
             self.stage = stage
 
     model = CurrentModel()
 
     # fit model
-    trainer = TrainerSubclass(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        checkpoint_callback=False
-    )
+    trainer = TrainerSubclass(default_root_dir=tmpdir, max_epochs=1, checkpoint_callback=False)
 
     trainer.fit(model)
     assert trainer.stage == 'fit'


Code review comments:
diff_hunk: @@ -471,13 +473,14 @@ def run_training_epoch(self):
 
         # run epoch
         for batch_idx, (batch, is_last_batch) in self.profiler.profile_iterable(
-                enumerate(_with_is_last(train_dataloader)), "get_train_batch"
+            enumerate(_with_is_last(train_dataloader)), "get_train_batch"
         ):
             # stop epoch if we limited the number of training batches
             if batch_idx >= self.num_training_batches:
                 break
 
             self.batch_idx = batch_idx
+            self.is_last_batch = is_last_batch
path: pytorch_lightning/trainer/training_loop.py
text: @user1:
let's pass this in as an argument to the function that needs it, like we do for "should_check_val" as example :)

@author:
@user1 if we pass `limit_train_batches` to 0.x or 1x, I think `is_last_batch` will still be `False` even if we reach the last batch of `num_training_batches`. Is my assumption correct?

@user1:
oh you're right. ahh, this variable has a misleading name. your code that had num_training_batches was better, sorry :)