Repository: iKernels/transformers-lightning (language: python)

Issue:
Electra scheduler and optimizer
Implementation of ELECTRA `scheduler` with `layerwise learning rate decay` and ELECTRA `optimizer` (a slightly modified version of AdamW). Original code provided in tensorflow here: https://github.com/google-research/electra/blob/master/model/optimization.py
refactored code, in particular the Metrics package to make it more co…
…mpatible with PL. It will be moved soon to that library. Fixing #63
Refactoring
Refactor code using the `yapf` code formatter.


Discussion / hints:

Fix patch:
diff --git a/.style.yapf b/.style.yapf
new file mode 100644
index 0000000..c1f5a5c
--- /dev/null
+++ b/.style.yapf
@@ -0,0 +1,10 @@
+[style]
+BASED_ON_STYLE = facebook
+SPACES_BEFORE_COMMENT = 4
+SPLIT_BEFORE_LOGICAL_OPERATOR = true
+ALLOW_MULTILINE_DICTIONARY_KEYS = true
+BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true
+BLANK_LINES_AROUND_TOP_LEVEL_DEFINITION = 2
+COLUMN_LIMIT = 120
+INDENT_WIDTH = 4
+SPLIT_COMPLEX_COMPREHENSION = true
diff --git a/.version.json b/.version.json
index eeb7920..5713036 100644
--- a/.version.json
+++ b/.version.json
@@ -1,5 +1,5 @@
 {
     "major": 0,
-    "minor": 4,
-    "patch": 5
+    "minor": 5,
+    "patch": 0
 }
diff --git a/README.md b/README.md
index 01ca3ab..63303e4 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,7 @@ A collection of `adapters`, `datasets`, `datamodules`, `callbacks`, `models`, `m
   * [2.7. Losses](transformers_lightning/losses)
   * [2.8. Language modeling](transformers_lightning/language_modeling)
   * [2.9. Schedulers](transformers_lightning/schedulers)
+  * [2.10. Optimizers](transformers_lightning/optimizers)
 
 **[3. Main file](#main)**
 
diff --git a/requirements.txt b/requirements.txt
index 5c2b312..c55baae 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,4 @@
 pytest
-pytorch-lightning >= 1.0.6
-transformers >= 3.5.0
+pytorch-lightning >= 1.1.2
+transformers >= 4.0.1
 torch >= 1.6.0
-scipy >= 1.5.4
diff --git a/setup.py b/setup.py
index 5c2685b..237304e 100644
--- a/setup.py
+++ b/setup.py
@@ -1,11 +1,13 @@
 import setuptools
 import json
 
+
 def load_long_description():
     with open("README.md", "r") as fh:
         long_description = fh.read()
     return long_description
 
+
 def load_requirements():
     requirements = []
     with open("requirements.txt", "r") as f:
@@ -15,12 +17,14 @@ def load_requirements():
                 requirements.append(line)
     return requirements
 
+
 def get_version():
     # get semver version [major.minor.patch]
     json_version = {}
     with open('.version.json', 'r') as f:
         json_version = json.load(f)
-    return '.'.join(str(w) for w in [json_version['major'],json_version['minor'],json_version['patch']])
+    return '.'.join(str(w) for w in [json_version['major'], json_version['minor'], json_version['patch']])
+
 
 setuptools.setup(
     name='transformers_lightning',
@@ -37,4 +41,4 @@ def get_version():
         "License :: OSI Approved :: GNU v2 License",
         "Operating System :: OS Independent",
     ]
-)
\ No newline at end of file
+)
diff --git a/transformers_lightning/adapters/file_adapter.py b/transformers_lightning/adapters/file_adapter.py
index 502a8cd..8e7af4b 100644
--- a/transformers_lightning/adapters/file_adapter.py
+++ b/transformers_lightning/adapters/file_adapter.py
@@ -14,11 +14,7 @@ def __init__(self, hparams: Namespace, filepath: str) -> None:
         :param filepath: path of the file that should be loaded
         """
         super().__init__(hparams)
-  
-        assert isinstance(filepath, str), (
-            f"Argument `filepath` must be of type `str`"
-        )
-        assert os.path.isfile(filepath), (
-            f"{filepath} is not a correct path to a file"
-        )
+
+        assert isinstance(filepath, str), (f"Argument `filepath` must be of type `str`")
+        assert os.path.isfile(filepath), (f"{filepath} is not a correct path to a file")
         self.filepath = filepath
diff --git a/transformers_lightning/adapters/transformers_adapter.py b/transformers_lightning/adapters/transformers_adapter.py
index 1686c2f..3317e51 100644
--- a/transformers_lightning/adapters/transformers_adapter.py
+++ b/transformers_lightning/adapters/transformers_adapter.py
@@ -25,7 +25,8 @@ class TransformersAdapter(CSVAdapter):
         represents tail tokens (tokens that starts with '##') and `False` represents first tokens in words.
     """
 
-    def __init__(self,
+    def __init__(
+        self,
         *args,
         tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast] = None,
         **kwargs
@@ -43,6 +44,4 @@ def _convert_ids_to_word_tails(self, ids: List[int]):
         >>> tokens = ['[CLS]', 'This', 'is', 'a', 's', '##yn', '##ta', '##ctic', '##ally', 'correct', 'sentence', '[SEP]']
         >>> words_tails = [False, False, False, False, False, True, True, True, True, False, False, False]
         """
-        return [
-            tok.startswith('##') for tok in self.tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=False)
-        ]
+        return [tok.startswith('##') for tok in self.tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=False)]
diff --git a/transformers_lightning/callbacks/transformers_model_checkpoint.py b/transformers_lightning/callbacks/transformers_model_checkpoint.py
index bb0d9ae..2ce614a 100644
--- a/transformers_lightning/callbacks/transformers_model_checkpoint.py
+++ b/transformers_lightning/callbacks/transformers_model_checkpoint.py
@@ -68,8 +68,8 @@ def on_train_start(self, trainer, pl_module):
         """ Check model can be saved and save hparams to understand what kind of experiment it was. """
         if trainer.global_rank != 0:
             return
-        
-        if not os.path.isdir(self.destination) :
+
+        if not os.path.isdir(self.destination):
             os.makedirs(self.destination)
 
         self.save_params()
@@ -87,16 +87,12 @@ def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, data
 
         # if global step is not multiple of checkpoint_interval
         if (
-            (self.hparams.checkpoint_interval is None) or
-            ((pl_module.global_step + 1) % self.hparams.checkpoint_interval) != 0
-        ):  
+            (self.hparams.checkpoint_interval is None)
+            or ((pl_module.global_step + 1) % self.hparams.checkpoint_interval) != 0
+        ):
             return
 
-        self.save_model(
-            pl_module,
-            epoch=trainer.current_epoch,
-            step=pl_module.global_step + 1
-        )
+        self.save_model(pl_module, epoch=trainer.current_epoch, step=pl_module.global_step + 1)
 
     @rank_zero_only
     def on_epoch_end(self, trainer, pl_module):
@@ -104,16 +100,12 @@ def on_epoch_end(self, trainer, pl_module):
         # only run on main process
         if trainer.global_rank != 0:
             return
-        
+
         # not validation checkpointing if it is disabled
         if self.hparams.no_epoch_checkpointing:
             return
 
-        self.save_model(
-            pl_module,
-            epoch=trainer.current_epoch,
-            step=pl_module.global_step + 1
-        )
+        self.save_model(pl_module, epoch=trainer.current_epoch, step=pl_module.global_step + 1)
 
     @rank_zero_only
     def on_train_end(self, trainer, pl_module):
@@ -125,12 +117,7 @@ def on_train_end(self, trainer, pl_module):
         if trainer.global_rank != 0:
             return
 
-        self.save_model(
-            pl_module,
-            epoch=trainer.current_epoch,
-            step=pl_module.global_step,
-            final=True
-        )
+        self.save_model(pl_module, epoch=trainer.current_epoch, step=pl_module.global_step, final=True)
 
     @rank_zero_only
     def on_validation_end(self, trainer, pl_module):
@@ -146,18 +133,26 @@ def on_validation_end(self, trainer, pl_module):
         if self.hparams.no_val_checkpointing:
             return
 
-        self.save_model(pl_module,
-            epoch=trainer.current_epoch,
-            step=pl_module.global_step + 1
-        )
+        self.save_model(pl_module, epoch=trainer.current_epoch, step=pl_module.global_step + 1)
 
     @staticmethod
     def add_callback_specific_args(parser):
         """ Add callback_specific arguments to parser. """
-        parser.add_argument('--checkpoint_interval', type=int, required=False, default=None,
-                            help="Save pre_trained models every steps. A None value means save only at the end of each epoch.")
-        parser.add_argument('--no_val_checkpointing', action="store_true",
-                            help="Disable transformers checkpointing at each validation end.")
-        parser.add_argument('--no_epoch_checkpointing', action="store_true",
-                            help="Disable transformers checkpointing at the end of each epoch.")
+        parser.add_argument(
+            '--checkpoint_interval',
+            type=int,
+            required=False,
+            default=None,
+            help="Save pre_trained models every steps. A None value means save only at the end of each epoch."
+        )
+        parser.add_argument(
+            '--no_val_checkpointing',
+            action="store_true",
+            help="Disable transformers checkpointing at each validation end."
+        )
+        parser.add_argument(
+            '--no_epoch_checkpointing',
+            action="store_true",
+            help="Disable transformers checkpointing at the end of each epoch."
+        )
         return parser
diff --git a/transformers_lightning/datamodules/__init__.py b/transformers_lightning/datamodules/__init__.py
index 54b71c3..018eb75 100644
--- a/transformers_lightning/datamodules/__init__.py
+++ b/transformers_lightning/datamodules/__init__.py
@@ -1,1 +1,1 @@
-from transformers_lightning.datamodules.superdatamodule import SuperDataModule
\ No newline at end of file
+from transformers_lightning.datamodules.superdatamodule import SuperDataModule
diff --git a/transformers_lightning/datamodules/superdatamodule.py b/transformers_lightning/datamodules/superdatamodule.py
index 23c5f61..8a57803 100644
--- a/transformers_lightning/datamodules/superdatamodule.py
+++ b/transformers_lightning/datamodules/superdatamodule.py
@@ -8,8 +8,7 @@
 import pytorch_lightning as pl
 from torch.utils.data import DataLoader
 from transformers_lightning import utils
-from transformers_lightning.datasets import (TransformersIterableDataset,
-                                             TransformersMapDataset)
+from transformers_lightning.datasets import (TransformersIterableDataset, TransformersMapDataset)
 
 
 class SuperDataModule(pl.LightningDataModule):
@@ -29,12 +28,8 @@ class SuperDataModule(pl.LightningDataModule):
     valid_adapter: SuperAdapter = None
     test_adapter: SuperAdapter = None
 
-    def __init__(self,
-        hparams,
-        train_adapter = None,
-        valid_adapter = None,
-        test_adapter = None,
-        collate_fn=utils.collate_single_fn
+    def __init__(
+        self, hparams, train_adapter=None, valid_adapter=None, test_adapter=None, collate_fn=utils.collate_single_fn
     ):
         super().__init__()
         self.hparams = hparams
@@ -56,11 +51,10 @@ def __init__(self,
 
             if isinstance(test_adapter, list):
                 for adapter in test_adapter:
-                    assert isinstance(adapter, SuperAdapter), (
-                        f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]"
-                    )
+                    assert isinstance(
+                        adapter, SuperAdapter
+                    ), (f"Argument `test_adapter` must be of type `SuperAdapter` or List[SuperAdapter]")
             self.test_adapter = test_adapter
-
         """
         This space should be used to instantiate the Adapters it they were not passed through the kwargs
 
@@ -83,16 +77,10 @@ def setup(self, stage=None):
 
         if stage == 'fit' or stage is None:
             if self.train_adapter is not None:
-                kwargs = {
-                    'start_from_step': self.hparams.skip_in_training
-                } if self.hparams.iterable_datasets else {}
-                self.train_dataset = dataset_class(
-                    self.hparams, self.train_adapter, self.trainer, **kwargs
-                )
+                kwargs = {'start_from_step': self.hparams.skip_in_training} if self.hparams.iterable_datasets else {}
+                self.train_dataset = dataset_class(self.hparams, self.train_adapter, self.trainer, **kwargs)
             if self.valid_adapter is not None:
-                self.valid_dataset = dataset_class(
-                    self.hparams, self.valid_adapter, self.trainer
-                )
+                self.valid_dataset = dataset_class(self.hparams, self.valid_adapter, self.trainer)
 
             assert self.train_adapter is None or self.train_dataset is not None, (
                 f"Cannot specify `train_adapter` and then `train_dataset` is None: "
@@ -106,13 +94,11 @@ def setup(self, stage=None):
         elif stage == 'test' or stage is None:
             if self.test_adapter is not None:
                 if isinstance(self.test_adapter, SuperAdapter):
-                    self.test_dataset = dataset_class(
-                        self.hparams, self.test_adapter, self.trainer
-                    )
+                    self.test_dataset = dataset_class(self.hparams, self.test_adapter, self.trainer)
                 else:
-                    self.test_dataset = [dataset_class(
-                        self.hparams, adapter, self.trainer
-                    ) for adapter in self.test_adapter]
+                    self.test_dataset = [
+                        dataset_class(self.hparams, adapter, self.trainer) for adapter in self.test_adapter
+                    ]
 
             assert self.test_adapter is None or self.test_dataset is not None, (
                 f"Cannot specify `test_adapter` and then `test_dataset` is None: "
@@ -130,12 +116,14 @@ def do_test(self):
 
     def default_dataloader(self, dataset, batch_size, shuffle=False):
         """ Return a dataloader with all usual default parameters. """
-        return DataLoader(dataset,
-                          batch_size=batch_size,
-                          num_workers=self.hparams.num_workers,
-                          pin_memory=True,
-                          collate_fn=self.collate_fn,
-                          shuffle=shuffle)
+        return DataLoader(
+            dataset,
+            batch_size=batch_size,
+            num_workers=self.hparams.num_workers,
+            pin_memory=True,
+            collate_fn=self.collate_fn,
+            shuffle=shuffle
+        )
 
     def train_dataloader(self):
         if self.train_dataset:
@@ -151,14 +139,22 @@ def val_dataloader(self):
     def test_dataloader(self):
         if self.test_adapter:
             if isinstance(self.test_dataset, list):
-                return [self.default_dataloader(dataset, self.hparams.test_batch_size, shuffle=False) for dataset in self.test_dataset]
+                return [
+                    self.default_dataloader(dataset, self.hparams.test_batch_size, shuffle=False)
+                    for dataset in self.test_dataset
+                ]
             return self.default_dataloader(self.test_dataset, self.hparams.test_batch_size, shuffle=False)
         return None
 
     @staticmethod
     def add_datamodule_specific_args(parser: ArgumentParser):
-        parser.add_argument('--num_workers', required=False, default=multiprocessing.cpu_count(), type=int,
-                            help='Number of workers to be used to load datasets')
+        parser.add_argument(
+            '--num_workers',
+            required=False,
+            default=multiprocessing.cpu_count(),
+            type=int,
+            help='Number of workers to be used to load datasets'
+        )
         parser.add_argument('--skip_in_training', type=int, default=None, required=False)
         parser.add_argument('--batch_size', type=int, default=32)
         parser.add_argument('--val_batch_size', type=int, default=256)
diff --git a/transformers_lightning/datasets/iterable_dataset.py b/transformers_lightning/datasets/iterable_dataset.py
index d927b6f..6b31cf7 100644
--- a/transformers_lightning/datasets/iterable_dataset.py
+++ b/transformers_lightning/datasets/iterable_dataset.py
@@ -70,9 +70,7 @@ def __init__(self, *args, start_from_step=None):
         """
         self.start_from_step = None
         if start_from_step is not None:
-            assert isinstance(start_from_step, int), (
-                f"`start_from` must be integer, found {start_from_step}"
-            )
+            assert isinstance(start_from_step, int), (f"`start_from` must be integer, found {start_from_step}")
 
             total_devices = utils.get_total_devices(trainer=self.trainer)
             effective_batch_size = self.hparams.batch_size * self.hparams.accumulate_grad_batches * total_devices
@@ -93,16 +91,12 @@ def __iter__(self):
         1) utils.batch_filter simply ensures that at least `world_size` elements are read at a time
         2) utils.filter_generator on distributed training to keep one element every `world_size`
         3) utils.filter_generator on parallel workers processing to keep one element every `num_workers`
-        """ 
+        """
         self.reader = iter(self.adapter)
-        
+
         # skip first steps if needed
         if self.start_from_step is not None:
-            self.reader = utils.filter_generator(
-                self.reader,
-                step=1,
-                offset=self.start_from_step
-            )
+            self.reader = utils.filter_generator(self.reader, step=1, offset=self.start_from_step)
             self.start_from_step = None
 
         # add distributed training logic
@@ -111,16 +105,9 @@ def __iter__(self):
             world_size = torch.distributed.get_world_size()
             rank = torch.distributed.get_rank()
 
-            self.reader = utils.batch_filter(
-                self.reader,
-                size=world_size
-            )
+            self.reader = utils.batch_filter(self.reader, size=world_size)
 
-            self.reader = utils.filter_generator(
-                self.reader,
-                step=world_size,
-                offset=rank
-            )
+            self.reader = utils.filter_generator(self.reader, step=world_size, offset=rank)
 
         # add parallel processing with workers logic
         if torch.utils.data.get_worker_info() is not None:
@@ -132,7 +119,7 @@ def __iter__(self):
 
         # pre-process data and return
         for line in self.reader:
-            
+
             preprocessed_data = self.adapter.preprocess_line(line)
 
             # if the adapter return a generator, yield an element per time
@@ -140,5 +127,3 @@ def __iter__(self):
                 yield from preprocessed_data
             else:
                 yield preprocessed_data
-
-
diff --git a/transformers_lightning/datasets/map_dataset.py b/transformers_lightning/datasets/map_dataset.py
index 8ffd17d..26b4be0 100644
--- a/transformers_lightning/datasets/map_dataset.py
+++ b/transformers_lightning/datasets/map_dataset.py
@@ -7,6 +7,7 @@ class TransformersMapDataset(SuperTransformersDataset, Dataset):
     Superclass of all map datasets. Tokenization is performed on the fly.
     Dataset is completely read into memory. If you need some lighter way of doing training, see `IterableDataset`
     """
+
     def __init__(self, *args):
         super().__init__(*args)
         self.data = list(iter(self.adapter))
@@ -21,9 +22,7 @@ def __getitem__(self, idx) -> dict:
                 raise ValueError("absolute value of index should not exceed dataset length")
             idx = len(self) + idx
 
-        assert 0 <= idx < len(self), (
-            f"Received index out of range {idx}, range: {0} <= idx < {len(self)}"
-        )
+        assert 0 <= idx < len(self), (f"Received index out of range {idx}, range: {0} <= idx < {len(self)}")
 
         row = self.data[idx]
         row_dict = self.adapter.preprocess_line(row)
diff --git a/transformers_lightning/datasets/stack_dataset.py b/transformers_lightning/datasets/stack_dataset.py
index 9cf8e63..698c74d 100644
--- a/transformers_lightning/datasets/stack_dataset.py
+++ b/transformers_lightning/datasets/stack_dataset.py
@@ -12,10 +12,10 @@ class StackDataset(torch.utils.data.Dataset):
 
     def __init__(self, *datasets):
         super().__init__()
-        
+
         assert len(datasets) > 0, 'datasets should not be an empty iterable'
         self.datasets = list(datasets)
-        
+
         if not utils.functional.all_equal_in_iterable([len(d) for d in self.datasets]):
             rank_zero_warn(
                 "Datasets do not have all the same length: "
@@ -28,9 +28,7 @@ def __len__(self):
     def __getitem__(self, idx):
         if idx < 0:
             if -idx > len(self):
-                raise ValueError(
-                    f"absolute value of index {idx} should not exceed dataset length {len(self)}"
-                )
+                raise ValueError(f"absolute value of index {idx} should not exceed dataset length {len(self)}")
             idx = len(self) + idx
         # assert index does not exit boundaries
         assert 0 <= idx < len(self), f"idx with value {idx} exists bounds [{0}, {len(self)}]"
diff --git a/transformers_lightning/datasets/superdataset.py b/transformers_lightning/datasets/superdataset.py
index 47801c9..f1e6c1a 100644
--- a/transformers_lightning/datasets/superdataset.py
+++ b/transformers_lightning/datasets/superdataset.py
@@ -15,7 +15,7 @@ class SuperTransformersDataset:
     def __init__(self, hparams: Namespace, adapter: SuperAdapter, trainer: Trainer):
         r"""
         Doing type check to avoid stupid errors related to arguments inversion
-        """        
+        """
         assert isinstance(hparams, Namespace), f"Argument `hparams` must be of type `Namespace`"
         assert isinstance(adapter, SuperAdapter), f"Argument `adapter` must be of type `SuperAdapter`"
         assert isinstance(trainer, Trainer), f"Argument `trainer` must be of type `Trainer`"
diff --git a/transformers_lightning/defaults/__init__.py b/transformers_lightning/defaults/__init__.py
index 3a18ae8..646f50b 100644
--- a/transformers_lightning/defaults/__init__.py
+++ b/transformers_lightning/defaults/__init__.py
@@ -1,4 +1,3 @@
-
 from argparse import ArgumentParser
 
 
@@ -17,12 +16,32 @@ class DefaultConfig:
 
     @staticmethod
     def add_defaults_args(parser: ArgumentParser):
-        parser.add_argument('--output_dir', type=str, required=False, default=DefaultConfig.output_dir,
-                            help='Specify a different output folder')
-        parser.add_argument('--pre_trained_dir', type=str, required=False, default=DefaultConfig.pre_trained_dir,
-                            help="Default path to save transformer models to")
-        parser.add_argument('--tensorboard_dir', type=str, required=False, default=DefaultConfig.tensorboard_dir,
-                            help="Where tensorboard logs should be saved")
-        parser.add_argument('--checkpoints_dir', type=str, required=False, default=DefaultConfig.checkpoints_dir,
-                            help="Where checkpoints should be saved")
-        return parser
\ No newline at end of file
+        parser.add_argument(
+            '--output_dir',
+            type=str,
+            required=False,
+            default=DefaultConfig.output_dir,
+            help='Specify a different output folder'
+        )
+        parser.add_argument(
+            '--pre_trained_dir',
+            type=str,
+            required=False,
+            default=DefaultConfig.pre_trained_dir,
+            help="Default path to save transformer models to"
+        )
+        parser.add_argument(
+            '--tensorboard_dir',
+            type=str,
+            required=False,
+            default=DefaultConfig.tensorboard_dir,
+            help="Where tensorboard logs should be saved"
+        )
+        parser.add_argument(
+            '--checkpoints_dir',
+            type=str,
+            required=False,
+            default=DefaultConfig.checkpoints_dir,
+            help="Where checkpoints should be saved"
+        )
+        return parser
diff --git a/transformers_lightning/language_modeling/masked_language_modeling.py b/transformers_lightning/language_modeling/masked_language_modeling.py
index 5978b4f..863e585 100644
--- a/transformers_lightning/language_modeling/masked_language_modeling.py
+++ b/transformers_lightning/language_modeling/masked_language_modeling.py
@@ -49,11 +49,9 @@ def __init__(
         self.mlm_probability = mlm_probability
         self.whole_word_masking = whole_word_masking
 
-    def __call__(
-        self,
-        inputs: torch.Tensor,        
-        words_tails: torch.Tensor = None
-    ) -> Tuple[torch.LongTensor, torch.LongTensor]:
+    def __call__(self,
+                 inputs: torch.Tensor,
+                 words_tails: torch.Tensor = None) -> Tuple[torch.LongTensor, torch.LongTensor]:
 
         if self.tokenizer.mask_token is None:
             raise ValueError(
@@ -61,10 +59,12 @@ def __call__(
             )
 
         device = inputs.device
-        labels = inputs.clone()      
+        labels = inputs.clone()
 
         # We sample a few tokens in each sequence for masked-LM training (with probability mlm_probability defaults to 0.15 in Bert/RoBERTa)
-        probability_matrix = torch.full(labels.shape, fill_value=self.mlm_probability, dtype=torch.float32, device=device)
+        probability_matrix = torch.full(
+            labels.shape, fill_value=self.mlm_probability, dtype=torch.float32, device=device
+        )
 
         # create whole work masking mask -> True if the token starts with ## (following token in composed words)
         if words_tails is None and self.whole_word_masking:
@@ -87,22 +87,21 @@ def __call__(
         # with whole word masking, assure all tokens in a word are either all masked or not
         if self.whole_word_masking:
             for i in range(1, masked_indices.shape[-1]):
-                masked_indices[:, i] = masked_indices[:, i] | (masked_indices[:, i-1] & words_tails[:, i])
+                masked_indices[:, i] = masked_indices[:, i] | (masked_indices[:, i - 1] & words_tails[:, i])
 
-        labels[~masked_indices] = IGNORE_IDX  # We only compute loss on masked tokens
+        labels[~masked_indices] = IGNORE_IDX    # We only compute loss on masked tokens
 
         # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
         indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8, device=device)).bool() & masked_indices
         inputs[indices_replaced] = self.tokenizer.mask_token_id
 
         # 10% of the time, we replace masked input tokens with random word
-        indices_random = torch.bernoulli(
-            torch.full(labels.shape, 0.5, device=device)
-        ).bool() & masked_indices & ~indices_replaced
+        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5, device=device)
+                                        ).bool() & masked_indices & ~indices_replaced
         random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long, device=device)
         inputs[indices_random] = random_words[indices_random]
 
         # The rest of the time (10% of the time) we keep the masked input tokens unchanged
         pass
-    
+
         return inputs, labels
diff --git a/transformers_lightning/language_modeling/random_token_substitution.py b/transformers_lightning/language_modeling/random_token_substitution.py
index e806eb2..6c2af0a 100644
--- a/transformers_lightning/language_modeling/random_token_substitution.py
+++ b/transformers_lightning/language_modeling/random_token_substitution.py
@@ -50,17 +50,17 @@ def __init__(
         self.rts_probability = rts_probability
         self.whole_word_swapping = whole_word_swapping
 
-    def __call__(
-        self,
-        inputs: torch.Tensor,
-        words_tails: torch.Tensor = None
-    ) -> Tuple[torch.LongTensor, torch.LongTensor]:
+    def __call__(self,
+                 inputs: torch.Tensor,
+                 words_tails: torch.Tensor = None) -> Tuple[torch.LongTensor, torch.LongTensor]:
 
         device = inputs.device
         labels = torch.full(inputs.shape, fill_value=0, dtype=torch.long, device=device)
 
         # We sample a few tokens in each sequence for masked-LM training (with probability args.rts_probability defaults to 0.15 in Bert/RoBERTa)
-        probability_matrix = torch.full(inputs.shape, fill_value=self.rts_probability, dtype=torch.float32, device=device)
+        probability_matrix = torch.full(
+            inputs.shape, fill_value=self.rts_probability, dtype=torch.float32, device=device
+        )
 
         # create whole work masking mask -> True if the token starts with ## (following token in composed words)
         if words_tails is None and self.whole_word_swapping:
@@ -89,7 +89,9 @@ def __call__(
         # with whole word masking, assure all tokens in a word are either all masked or not
         if self.whole_word_swapping:
             for i in range(1, substituted_indices.shape[-1]):
-                substituted_indices[:, i] = substituted_indices[:, i] | (substituted_indices[:, i-1] & words_tails[:, i])
+                substituted_indices[:,
+                                    i] = substituted_indices[:,
+                                                             i] | (substituted_indices[:, i - 1] & words_tails[:, i])
 
         random_words = torch.randint(len(self.tokenizer), inputs.shape, dtype=torch.long, device=device)
         inputs[substituted_indices] = random_words[substituted_indices]
diff --git a/transformers_lightning/language_modeling/utils.py b/transformers_lightning/language_modeling/utils.py
index f33dc0a..be3adcd 100644
--- a/transformers_lightning/language_modeling/utils.py
+++ b/transformers_lightning/language_modeling/utils.py
@@ -3,13 +3,18 @@
 from transformers import PreTrainedTokenizer
 
 
-def whole_word_tails_mask(inputs: torch.Tensor, tokenizer: PreTrainedTokenizer, device: torch.device = None) -> torch.BoolTensor:
-   r"""
+def whole_word_tails_mask(
+    inputs: torch.Tensor, tokenizer: PreTrainedTokenizer, device: torch.device = None
+) -> torch.BoolTensor:
+    r"""
    # create whole work masking mask -> 1 if the token starts with ## (following token in composed words)
    """
-   return torch.tensor(
-      [[token.startswith('##') for token in tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=False)]
-      for ids in inputs],
-      device=device,
-      dtype=torch.bool
-   )
+    return torch.tensor(
+        [
+            [token.startswith('##')
+             for token in tokenizer.convert_ids_to_tokens(ids, skip_special_tokens=False)]
+            for ids in inputs
+        ],
+        device=device,
+        dtype=torch.bool
+    )
diff --git a/transformers_lightning/metrics/README.md b/transformers_lightning/metrics/README.md
index 8f42415..ec6432f 100644
--- a/transformers_lightning/metrics/README.md
+++ b/transformers_lightning/metrics/README.md
@@ -16,15 +16,15 @@ Since information retrieval metric must work on many predictions at a time, we i
 Compute the MRR score over a set of documents.
 
 Usage example:
+
 ```python
 >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
 >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
 >>> target = torch.tensor([False, False, True, False, True, False, False])
-
 >>> mrr = MeanReciprocalRank()
 >>> mrr(indexes, preds, target)
 >>> mrr.compute()
-... 0.75
+tensor(0.75)
 ```
 
 ### MeanAveragePrecision
@@ -32,67 +32,66 @@ Usage example:
 Compute the MAP score over a set of documents.
 
 Usage example:
+
 ```python
 >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
 >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
 >>> target = torch.tensor([False, False, True, False, True, False, False])
-
 >>> map = MeanAveragePrecision()
 >>> map(indexes, preds, target)
 >>> map.compute()
-... 0.75
+tensor(0.75)
 ```
 
 
-### PrecisionAtK
+### Precision
 
-Compute the Precision @ K over a set of documents.
+Compute the Precision (possibly @ K) over a set of documents.
 
 Usage example:
+
 ```python
 >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
 >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
 >>> target = torch.tensor([False, False, True, False, True, False, False])
-
->>> p_k = PrecisionAtK(k=1)
+>>> p_k = Precision(k=1)
 >>> p_k(indexes, preds, target)
 >>> p_k.compute()
-... 0.5
+tensor(0.5)
 ```
 
 
-### RecallAtK
+### Recall
 
-Compute the Recall @ K over a set of documents.
+Compute the Recall (possibly @ K) over a set of documents.
 
 Usage example:
+
 ```python
 >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
 >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
 >>> target = torch.tensor([False, False, True, False, True, False, False])
-
->>> r_k = RecallAtK(k=1)
+>>> r_k = Recall(k=1)
 >>> r_k(indexes, preds, target)
 >>> r_k.compute()
-... 0.5
+tensor(0.5)
 ```
 
 
-
 ### HitRateAtK
 
 Compute the Hit Rate @ K over a set of documents.
 
 Usage example:
+
 ```python
 >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
 >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
 >>> target = torch.tensor([False, False, True, False, True, False, False])
-
 >>> hr_k = HitRateAtK(k=1)
 >>> hr_k(indexes, preds, target)
 >>> hr_k.compute()
-... 0.5
+tensor(0.5)
 ```
 
 
diff --git a/transformers_lightning/metrics/functional/average_precision.py b/transformers_lightning/metrics/functional/average_precision.py
index d99331a..eb0168d 100644
--- a/transformers_lightning/metrics/functional/average_precision.py
+++ b/transformers_lightning/metrics/functional/average_precision.py
@@ -1,40 +1,40 @@
 import torch
 
 
-def average_precision(
-    preds: torch.Tensor,
-    target: torch.Tensor
-):
+def average_precision(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
     r"""
-    Computes average precision metric for information retrieval,
-    as explained here: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Mean_average_precision
+    Computes average precision (for information retrieval), as explained
+    `here <https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision>`_.
 
-    `preds` and `target` should be of the same shape and live on the same device. If not target is true, 0 is returned.
+    `preds` and `target` should be of the same shape and live on the same device. If no `target` is ``True``,
+    0 is returned. Target must be of type `bool` or `int`, otherwise an error is raised.
 
     Args:
         preds: estimated probabilities of each document to be relevant.
-        target: ground truth about each document being relevant.
-    
-    Returns:
+        target: ground truth about each document being relevant or not. Requires `bool` or `int` tensor.
+
+    Return:
         a single-value tensor with the average precision (AP) of the predictions `preds` wrt the labels `target`.
 
     Example:
-
         >>> preds = torch.tensor([0.2, 0.3, 0.5])
         >>> target = torch.tensor([True, False, True])
         >>> average_precision(preds, target)
-        ... 0.833
+        tensor(0.8333)
     """
 
-    if preds.shape != target.shape or preds.device != target.device: 
-        raise ValueError(
-            f"`preds` and `target` must have the same shape and be on the same device"
-        )
+    if preds.shape != target.shape or preds.device != target.device:
+        raise ValueError("`preds` and `target` must have the same shape and live on the same device")
+
+    if target.dtype not in (torch.bool, torch.int16, torch.int32, torch.int64):
+        raise ValueError("`target` must be a tensor of booleans or integers")
+
+    if target.dtype is not torch.bool:
+        target = target.bool()
 
     if target.sum() == 0:
-        return torch.tensor([0]).to(preds)
+        return torch.tensor(0, device=preds.device)
 
     target = target[torch.argsort(preds, dim=-1, descending=True)]
-    positions = torch.arange(1, len(target) + 1, device=target.device)[target > 0]
-    res = torch.true_divide((torch.arange(len(positions), device=positions.device) + 1), positions).mean()
-    return res
+    positions = torch.arange(1, len(target) + 1, device=target.device, dtype=torch.float32)[target > 0]
+    return torch.div((torch.arange(len(positions), device=positions.device, dtype=torch.float32) + 1), positions).mean()
diff --git a/transformers_lightning/metrics/functional/hit_rate.py b/transformers_lightning/metrics/functional/hit_rate.py
index 40484f7..0602391 100644
--- a/transformers_lightning/metrics/functional/hit_rate.py
+++ b/transformers_lightning/metrics/functional/hit_rate.py
@@ -1,37 +1,40 @@
 import torch
 
 
-def hit_rate(
-    preds: torch.Tensor,
-    target: torch.Tensor,
-    k: int = 1
-):
+def hit_rate(preds: torch.Tensor, target: torch.Tensor, k: int = 1):
     r"""
-    Computes the hit_rate @ k metric for information retrieval
-    Hir Rate at k is 1 iff there is at least one relevant documents among the top K.
+    Computes hit rate (for information retrieval).
+    Hir Rate is 1.0 iff there is at least one relevant documents among the top `k`.
 
-    `preds` and `target` should be of the same shape and live on the same device. If not target is true, 0 is returned.
+    `preds` and `target` should be of the same shape and live on the same device. If no `target` is ``True``,
+    0 is returned. Target must be of type `bool` or `int`, otherwise an error is raised.
 
     Args:
         preds: estimated probabilities of each document to be relevant.
-        target: ground truth about each document being relevant.
+        target: ground truth about each document being relevant or not. Requires `bool` or `int` tensor.
         k: consider only the top k elements.
 
-    Returns:
-        a single-value tensor with the hit rate at k (HR@K) of the predictions `preds` wrt the labels `target`.
+    Return:
+        a single-value tensor with the hit rate (HR) of the predictions `preds` wrt the labels `target`.
 
     Example:
-
         >>> preds = torch.tensor([0.2, 0.3, 0.5])
         >>> target = torch.tensor([True, False, True])
-        >>> hit_rate(preds, target, k=2)
-        ... 1.0
+        >>> hit_rate(preds, target)
+        tensor(1.0)
     """
 
-    if preds.shape != target.shape or preds.device != target.device: 
-        raise ValueError(
-            f"`preds` and `target` must have the same shape and be on the same device"
-        )
+    if k <= 0:
+        raise ValueError("`k` must be an integer greater or equal to `0`")
+
+    if preds.shape != target.shape or preds.device != target.device:
+        raise ValueError("`preds` and `target` must have the same shape and live on the same device")
+
+    if target.dtype not in (torch.bool, torch.int16, torch.int32, torch.int64):
+        raise ValueError("`target` must be a tensor of booleans or integers")
+
+    if target.dtype is not torch.bool:
+        target = target.bool()
 
     relevant = target[torch.argsort(preds, dim=-1, descending=True)][:k].sum()
-    return (relevant > 0).to(preds)
+    return (relevant > 0).float()
diff --git a/transformers_lightning/metrics/functional/ir_precision.py b/transformers_lightning/metrics/functional/ir_precision.py
index de020c6..8557980 100644
--- a/transformers_lightning/metrics/functional/ir_precision.py
+++ b/transformers_lightning/metrics/functional/ir_precision.py
@@ -1,41 +1,41 @@
 import torch
 
 
-def precision(
-    preds: torch.Tensor,
-    target: torch.Tensor,
-    k: int = 1
-):
+def precision(preds: torch.Tensor, target: torch.Tensor, k: int = None):
     r"""
-    Computes the precision @ k metric for information retrieval,
-    as explained here: https://en.wikipedia.org/wiki/Precision_and_recall#Definition_(information_retrieval_context)
+    Computes the precision metric (for information retrieval),
+    as explained `here <https://en.wikipedia.org/wiki/Precision_and_recall#Precision>`_.
     Precision at K is the fraction of relevant documents among the top K.
 
-    `preds` and `target` should be of the same shape and live on the same device. If not target is true, 0 is returned.
+    `preds` and `target` should be of the same shape and live on the same device. If no `target` is ``True``,
+    0 is returned. Target must be of type `bool` or `int`, otherwise an error is raised.
 
     Args:
         preds: estimated probabilities of each document to be relevant.
-        target: ground truth about each document being relevant.
+        target: ground truth about each document being relevant or not. Requires `bool` or `int` tensor.
         k: consider only the top k elements.
 
     Returns:
         a single-value tensor with the precision at k (P@K) of the predictions `preds` wrt the labels `target`.
 
     Example:
-
         >>> preds = torch.tensor([0.2, 0.3, 0.5])
         >>> target = torch.tensor([True, False, True])
         >>> precision(preds, target, k=2)
-        ... 0.5
+        ... tensor(0.5)
     """
 
-    if preds.shape != target.shape or preds.device != target.device: 
-        raise ValueError(
-            f"`preds` and `target` must have the same shape and be on the same device"
-        )
+    if preds.shape != target.shape or preds.device != target.device:
+        raise ValueError("`preds` and `target` must have the same shape and live on the same device")
+
+    if target.dtype not in (torch.bool, torch.int16, torch.int32, torch.int64):
+        raise ValueError("`target` must be a tensor of booleans or integers")
+
+    if target.dtype is not torch.bool:
+        target = target.bool()
 
     if target.sum() == 0:
-        return torch.tensor([0]).to(preds)
+        return torch.tensor(0, device=preds.device)
 
     relevant = target[torch.argsort(preds, dim=-1, descending=True)][:k].sum()
     return torch.true_divide(relevant, k)
diff --git a/transformers_lightning/metrics/functional/ir_recall.py b/transformers_lightning/metrics/functional/ir_recall.py
index b26af26..de0526d 100644
--- a/transformers_lightning/metrics/functional/ir_recall.py
+++ b/transformers_lightning/metrics/functional/ir_recall.py
@@ -1,41 +1,41 @@
 import torch
 
 
-def recall(
-    preds: torch.Tensor,
-    target: torch.Tensor,
-    k: int = 1
-):
+def recall(preds: torch.Tensor, target: torch.Tensor, k: int = None):
     r"""
-    Computes the recall @ k metric for information retrieval,
-    as explained here: https://en.wikipedia.org/wiki/Precision_and_recall#Definition_(information_retrieval_context)
+    Computes the recall metric (for information retrieval),
+    as explained `here <https://en.wikipedia.org/wiki/Precision_and_recall#Recall>`_.
     Recall at K is the fraction of relevant documents in top K among all the relevant documents.
 
-    `preds` and `target` should be of the same shape and live on the same device. If not target is true, 0 is returned.
+    `preds` and `target` should be of the same shape and live on the same device. If no `target` is ``True``,
+    0 is returned. Target must be of type `bool` or `int`, otherwise an error is raised.
 
     Args:
         preds: estimated probabilities of each document to be relevant.
-        target: ground truth about each document being relevant.
+        target: ground truth about each document being relevant or not. Requires `bool` or `int` tensor.
         k: consider only the top k elements.
 
     Returns:
         a single-value tensor with the recall at k (R@K) of the predictions `preds` wrt the labels `target`.
 
     Example:
-
         >>> preds = torch.tensor([0.2, 0.3, 0.5])
         >>> target = torch.tensor([True, False, True])
         >>> recall(preds, target, k=2)
-        ... 0.5
+        ... tensor(0.5)
     """
 
-    if preds.shape != target.shape or preds.device != target.device: 
-        raise ValueError(
-            f"`preds` and `target` must have the same shape and be on the same device"
-        )
+    if preds.shape != target.shape or preds.device != target.device:
+        raise ValueError("`preds` and `target` must have the same shape and live on the same device")
+
+    if target.dtype not in (torch.bool, torch.int16, torch.int32, torch.int64):
+        raise ValueError("`target` must be a tensor of booleans or integers")
+
+    if target.dtype is not torch.bool:
+        target = target.bool()
 
     if target.sum() == 0:
-        return torch.tensor([0]).to(preds)
+        return torch.tensor(0, device=preds.device)
 
     relevant = target[torch.argsort(preds, dim=-1, descending=True)][:k].sum()
     return torch.true_divide(relevant, target.sum())
diff --git a/transformers_lightning/metrics/functional/reciprocal_rank.py b/transformers_lightning/metrics/functional/reciprocal_rank.py
index 9976e65..88e2aed 100644
--- a/transformers_lightning/metrics/functional/reciprocal_rank.py
+++ b/transformers_lightning/metrics/functional/reciprocal_rank.py
@@ -1,38 +1,39 @@
 import torch
 
 
-def reciprocal_rank(
-    preds: torch.Tensor,
-    target: torch.Tensor
-):
+def reciprocal_rank(preds: torch.Tensor, target: torch.Tensor):
     r"""
-    Computes reciprocal rank metric for information retrieval,
-    as explained here: https://en.wikipedia.org/wiki/Mean_reciprocal_rank
+    Computes reciprocal rank (for information retrieval), as explained
+    `here <https://en.wikipedia.org/wiki/Mean_reciprocal_rank>`.
 
-    `preds` and `target` should be of the same shape and live on the same device. If not target is true, 0 is returned.
+    `preds` and `target` should be of the same shape and live on the same device. If no `target` is ``True``,
+    0 is returned. Target must be of type `bool` or `int`, otherwise an error is raised.
 
     Args:
         preds: estimated probabilities of each document to be relevant.
-        target: ground truth about each document being relevant.
+        target: ground truth about each document being relevant or not. Requires `bool` or `int` tensor.
 
     Returns:
         a single-value tensor with the reciprocal rank (RR) of the predictions `preds` wrt the labels `target`.
 
     Example:
-
         >>> preds = torch.tensor([0.2, 0.3, 0.5])
         >>> target = torch.tensor([False, True, False])
         >>> reciprocal_rank(preds, target)
-        ... 0.5
+        ... tensor(0.5)
     """
 
-    if preds.shape != target.shape or preds.device != target.device: 
-        raise ValueError(
-            f"`preds` and `target` must have the same shape and be on the same device"
-        )
+    if preds.shape != target.shape or preds.device != target.device:
+        raise ValueError("`preds` and `target` must have the same shape and live on the same device")
+
+    if target.dtype not in (torch.bool, torch.int16, torch.int32, torch.int64):
+        raise ValueError("`target` must be a tensor of booleans or integers")
+
+    if target.dtype is not torch.bool:
+        target = target.bool()
 
     if target.sum() == 0:
-        return torch.tensor(0).to(preds)
+        return torch.tensor(0, device=preds.device)
 
     target = target[torch.argsort(preds, dim=-1, descending=True)]
     position = torch.where(target == 1)[0]
diff --git a/transformers_lightning/metrics/retrieval/__init__.py b/transformers_lightning/metrics/retrieval/__init__.py
index 350bd95..bb589ac 100644
--- a/transformers_lightning/metrics/retrieval/__init__.py
+++ b/transformers_lightning/metrics/retrieval/__init__.py
@@ -1,5 +1,5 @@
 from transformers_lightning.metrics.retrieval.mean_average_precision import MeanAveragePrecision
 from transformers_lightning.metrics.retrieval.mean_reciprocal_rank import MeanReciprocalRank
-from transformers_lightning.metrics.retrieval.precision import PrecisionAtK
-from transformers_lightning.metrics.retrieval.recall import RecallAtK
+from transformers_lightning.metrics.retrieval.precision import Precision
+from transformers_lightning.metrics.retrieval.recall import Recall
 from transformers_lightning.metrics.retrieval.hit_rate import HitRateAtK
diff --git a/transformers_lightning/metrics/retrieval/hit_rate.py b/transformers_lightning/metrics/retrieval/hit_rate.py
index da121e3..ad40c28 100644
--- a/transformers_lightning/metrics/retrieval/hit_rate.py
+++ b/transformers_lightning/metrics/retrieval/hit_rate.py
@@ -1,4 +1,4 @@
-from typing import List
+import torch
 
 from transformers_lightning.metrics.retrieval.retrieval_metric import RetrievalMetric
 from transformers_lightning.metrics.functional.hit_rate import hit_rate
@@ -6,30 +6,60 @@
 
 class HitRateAtK(RetrievalMetric):
     r"""
-    Hit Rate at K computes the HR@K over multiple retrieved documents for each query.
-    Each hit rate at k computation over a single query can be done on a different number 
-    of predictions thanks to the usage of a tensor dedicated to separate query results.
+    Computes `Hit Rate`.
+    Works with binary data. Accepts integer or float predictions from a model output.
 
-    Notice that HR@1 == P@1
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+    `indexes`, `preds` and `target` must have the same dimension.
+    `indexes` indicate to which query a prediction belongs.
 
-    Example:
+    Predictions will be first grouped by indexes and then HR will be computed as the mean
+    of the HR over each query.
+
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
 
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+
+    Example:
         >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
         >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
         >>> target = torch.tensor([False, False, True, False, True, False, False])
-
         >>> hr_k = HitRateAtK(k=1)
         >>> hr_k(indexes, preds, target)
+        tensor(0.5)
         >>> hr_k.compute()
-        ... 0.5
+        tensor(0.5)
     """
 
-    def __init__(self, *args, k=1, **kwargs):
+    def __init__(self, *args, k: int = 1, **kwargs):
         super().__init__(*args, **kwargs)
+
+        if not k or k <= 0:
+            raise ValueError("`k` must be an integer greater than 0")
+
         self.k = k
 
-    def metric(self, group: List[int]):
-        _preds = self.preds[group]
-        _target = self.target[group]
-        valid_indexes = (_target != self.exclude)
-        return hit_rate(_preds[valid_indexes], _target[valid_indexes], k=self.k)
\ No newline at end of file
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        valid_indexes = target != self.exclude
+        return hit_rate(preds[valid_indexes], target[valid_indexes], k=self.k)
diff --git a/transformers_lightning/metrics/retrieval/mean_average_precision.py b/transformers_lightning/metrics/retrieval/mean_average_precision.py
index d036d17..97c2d40 100644
--- a/transformers_lightning/metrics/retrieval/mean_average_precision.py
+++ b/transformers_lightning/metrics/retrieval/mean_average_precision.py
@@ -1,4 +1,4 @@
-from typing import List
+import torch
 
 from transformers_lightning.metrics.retrieval.retrieval_metric import RetrievalMetric
 from transformers_lightning.metrics.functional.average_precision import average_precision
@@ -6,24 +6,53 @@
 
 class MeanAveragePrecision(RetrievalMetric):
     r"""
-    Mean Average Precision computes the MAP over multiple retrieved documents for each query.
-    Each average precision computation on a single query can be done on a different number of 
-    predictions thanks to the usage of a tensor dedicated to separate query results.
+    Computes `Mean Average Precision
+    <https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Mean_average_precision>`_.
+    Works with binary data. Accepts integer or float predictions from a model output.
 
-    Example:
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+    `indexes`, `preds` and `target` must have the same dimension.
+    `indexes` indicate to which query a prediction belongs.
+
+    Predictions will be first grouped by indexes and then MAP will be computed as the mean
+    of the Average Precisions over each query.
+
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
 
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+
+    Example:
         >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
         >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
         >>> target = torch.tensor([False, False, True, False, True, False, False])
-
         >>> map = MeanAveragePrecision()
         >>> map(indexes, preds, target)
+        tensor(0.7500)
         >>> map.compute()
-        ... 0.75
+        tensor(0.7500)
     """
 
-    def metric(self, group: List[int]):
-        _preds = self.preds[group]
-        _target = self.target[group]
-        valid_indexes = _target != self.exclude
-        return average_precision(_preds[valid_indexes], _target[valid_indexes])
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        valid_indexes = target != self.exclude
+        return average_precision(preds[valid_indexes], target[valid_indexes])
diff --git a/transformers_lightning/metrics/retrieval/mean_reciprocal_rank.py b/transformers_lightning/metrics/retrieval/mean_reciprocal_rank.py
index 334b987..e84810b 100644
--- a/transformers_lightning/metrics/retrieval/mean_reciprocal_rank.py
+++ b/transformers_lightning/metrics/retrieval/mean_reciprocal_rank.py
@@ -1,4 +1,4 @@
-from typing import List
+import torch
 
 from transformers_lightning.metrics.retrieval.retrieval_metric import RetrievalMetric
 from transformers_lightning.metrics.functional.reciprocal_rank import reciprocal_rank
@@ -6,24 +6,51 @@
 
 class MeanReciprocalRank(RetrievalMetric):
     r"""
-    Mean Reciprocal Rank computes the MRR over multiple retrieved documents for each query.
-    Each reciprocal rank computation on a single query can be done on a different number of
-    predictions thanks to the usage of a tensor dedicated to separate query results.
+    Computes `Mean Reciprocal Rank <https://en.wikipedia.org/wiki/Mean_reciprocal_rank>`_.
+    Works with binary data. Accepts integer or float predictions from a model output.
 
-    Example:
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+    `indexes`, `preds` and `target` must have the same dimension.
+    `indexes` indicate to which query a prediction belongs.
+
+    Predictions will be first grouped by indexes and then MRR will be computed as the mean
+    of the Reciprocal Rank over each query.
+
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
 
+    Example:
         >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
         >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
         >>> target = torch.tensor([False, False, True, False, True, False, False])
-
-        >>> mrr = MeanReciprocalRank()
-        >>> mrr(indexes, preds, target)
-        >>> mrr.compute()
-        ... 0.75
+        >>> map = MeanReciprocalRank()
+        >>> map(indexes, preds, target)
+        tensor(0.7500)
+        >>> map.compute()
+        tensor(0.7500)
     """
 
-    def metric(self, group: List[int]):
-        _preds = self.preds[group]
-        _target = self.target[group]
-        valid_indexes = (_target != self.exclude)
-        return reciprocal_rank(_preds[valid_indexes], _target[valid_indexes])
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        valid_indexes = target != self.exclude
+        return reciprocal_rank(preds[valid_indexes], target[valid_indexes])
diff --git a/transformers_lightning/metrics/retrieval/precision.py b/transformers_lightning/metrics/retrieval/precision.py
index b7e856e..7d0d5e7 100644
--- a/transformers_lightning/metrics/retrieval/precision.py
+++ b/transformers_lightning/metrics/retrieval/precision.py
@@ -1,33 +1,62 @@
-from typing import List
+import torch
 
 from transformers_lightning.metrics.retrieval.retrieval_metric import RetrievalMetric
 from transformers_lightning.metrics.functional.ir_precision import precision
 
 
-class PrecisionAtK(RetrievalMetric):
+class Precision(RetrievalMetric):
     r"""
-    Precision at K computes the P@K over multiple retrieved documents for each query.
-    Each precision at k computation on each query can be done on a different number
-    of predictions thanks to the usage of a tensor dedicated to separate query results.
+    Computes `Precision
+    <https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Precision>`_.
+    Works with binary data. Accepts integer or float predictions from a model output.
+
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+    `indexes`, `preds` and `target` must have the same dimension.
+    `indexes` indicate to which query a prediction belongs.
+
+    Predictions will be first grouped by indexes and then Precision will be computed as the mean
+    of the Precision over each query.
+
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
+
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
 
     Example:
-
         >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
         >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
         >>> target = torch.tensor([False, False, True, False, True, False, False])
-
-        >>> p_k = PrecisionAtK(k=1)
+        >>> p_k = Precision(k=1)
         >>> p_k(indexes, preds, target)
+        tensor(0.5)
         >>> p_k.compute()
-        ... 0.5
+        tensor(0.5)
     """
 
-    def __init__(self, *args, k=1, **kwargs):
+    def __init__(self, *args, k: int = 1, **kwargs):
         super().__init__(*args, **kwargs)
         self.k = k
 
-    def metric(self, group: List[int]):
-        _preds = self.preds[group]
-        _target = self.target[group]
-        valid_indexes = (_target != self.exclude)
-        return precision(_preds[valid_indexes], _target[valid_indexes], k=self.k)
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        valid_indexes = target != self.exclude
+        return precision(preds[valid_indexes], target[valid_indexes], k=self.k)
diff --git a/transformers_lightning/metrics/retrieval/recall.py b/transformers_lightning/metrics/retrieval/recall.py
index 0e84413..8fb159f 100644
--- a/transformers_lightning/metrics/retrieval/recall.py
+++ b/transformers_lightning/metrics/retrieval/recall.py
@@ -1,33 +1,62 @@
-from typing import List
+import torch
 
 from transformers_lightning.metrics.retrieval.retrieval_metric import RetrievalMetric
 from transformers_lightning.metrics.functional.ir_recall import recall
 
 
-class RecallAtK(RetrievalMetric):
+class Recall(RetrievalMetric):
     r"""
-    Recall at K computes the R@K over multiple retrieved documents for each query.
-    Each recall at k computation on a single query can be done on a different number
-    of predictions thanks to the usage of a tensor dedicated to separate query results.
+    Computes `Recall
+    <https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Recall>`_.
+    Works with binary data. Accepts integer or float predictions from a model output.
+
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+    `indexes`, `preds` and `target` must have the same dimension.
+    `indexes` indicate to which query a prediction belongs.
+
+    Predictions will be first grouped by indexes and then Recall will be computed as the mean
+    of the Recall over each query.
+
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
+
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
 
     Example:
-
         >>> indexes = torch.tensor([0, 0, 0, 1, 1, 1, 1])
         >>> preds = torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
         >>> target = torch.tensor([False, False, True, False, True, False, False])
-
-        >>> r_k = RecallAtK(k=1)
-        >>> r_k(indexes, preds, target)
-        >>> r_k.compute()
-        ... 0.5
+        >>> map = Recall(k=1)
+        >>> map(indexes, preds, target)
+        tensor(0.5)
+        >>> map.compute()
+        tensor(0.5)
     """
 
-    def __init__(self, *args, k=1, **kwargs):
+    def __init__(self, *args, k: int = None, **kwargs):
         super().__init__(*args, **kwargs)
         self.k = k
 
-    def metric(self, group: List[int]):
-        _preds = self.preds[group]
-        _target = self.target[group]
-        valid_indexes = (_target != self.exclude)
-        return recall(_preds[valid_indexes], _target[valid_indexes], k=self.k)
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        valid_indexes = target != self.exclude
+        return recall(preds[valid_indexes], target[valid_indexes], k=self.k)
diff --git a/transformers_lightning/metrics/retrieval/retrieval_metric.py b/transformers_lightning/metrics/retrieval/retrieval_metric.py
index da17b62..6eb6663 100644
--- a/transformers_lightning/metrics/retrieval/retrieval_metric.py
+++ b/transformers_lightning/metrics/retrieval/retrieval_metric.py
@@ -1,72 +1,137 @@
 import torch
-from typing import List
+from abc import abstractmethod, ABC
+from typing import Callable, Optional, Any
 from pytorch_lightning.metrics import Metric
 
 from transformers_lightning.metrics.utils import get_mini_groups
 
-
 IGNORE_IDX = -100
 
-class RetrievalMetric(Metric):
+
+class RetrievalMetric(Metric, ABC):
     r"""
-    Compute a metric for Information Retrieval by grouping predictions on the same
-    document using indexes. Detailed information about metrics are contained
-    in sub-classes.
-
-    It may be possible that a document has no positive label: this case can
-    be managed in different ways using the `empty_documents` parameter:
-    - `skip`: those documents are skipped (default)
-    - `error`: a `ValueError` is raised
-    - `positive`: those documents are counted as positive predictions (1.0)
-    - `negative`: those documents are counted as negative predictions (0.0)
-
-    Entries with targets equal to `exclude` will be ignored.
-    Subclasses must override at least the `metric` method.
-    """
+    Works with binary data. Accepts integer or float predictions from a model output.
+
+    Forward accepts
+    - ``indexes`` (long tensor): ``(N, ...)``
+    - ``preds`` (float or int tensor): ``(N, ...)``
+    - ``target`` (long or bool tensor): ``(N, ...)``
+
+    `indexes`, `preds` and `target` must have the same dimension and will be flatten
+    to single dimension once provided.
 
-    options = ['error', 'skip', 'positive', 'negative']
+    `indexes` indicate to which query a prediction belongs.
+    Predictions will be first grouped by indexes. Then the
+    real metric, defined by overriding the `_metric` method,
+    will be computed as the mean of the scores over each query.
 
-    def __init__(self, dist_sync_on_step: bool = False, empty_documents: str = 'skip', exclude: int = IGNORE_IDX):
-        super().__init__(dist_sync_on_step=dist_sync_on_step)
+    Args:
+        query_without_relevant_docs:
+            Specify what to do with queries that do not have at least a positive target. Choose from:
+        exclude:
+            Do not take into account predictions where the target is equal to this value. default `-100`
+        compute_on_step:
+            Forward only calls ``update()`` and return None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step. default: False
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects
+            the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+
+            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned
+            - ``'error'``: raise a ``ValueError``
+            - ``'pos'``: score on those queries is counted as ``1.0``
+            - ``'neg'``: score on those queries is counted as ``0.0``
+
+    """
+
+    def __init__(
+        self,
+        query_without_relevant_docs: str = 'skip',
+        exclude: int = IGNORE_IDX,
+        compute_on_step: bool = True,
+        dist_sync_on_step: bool = False,
+        process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None
+    ):
+        super().__init__(
+            compute_on_step=compute_on_step,
+            dist_sync_on_step=dist_sync_on_step,
+            process_group=process_group,
+            dist_sync_fn=dist_sync_fn
+        )
 
-        if empty_documents not in self.options:
+        query_without_relevant_docs_options = ('error', 'skip', 'pos', 'neg')
+        if query_without_relevant_docs not in query_without_relevant_docs_options:
             raise ValueError(
-                f"`empty_documents` received a wrong value {empty_documents}."
-                f"Allowed values are {self.options}"
+                f"`query_without_relevant_docs` received a wrong value {query_without_relevant_docs}. "
+                f"Allowed values are {query_without_relevant_docs_options}"
             )
 
-        self.empty_documents = empty_documents
+        self.query_without_relevant_docs = query_without_relevant_docs
         self.exclude = exclude
 
-        self.add_state("idx", default=torch.tensor([], dtype=torch.int64), dist_reduce_fx="cat")
-        self.add_state("preds", default=torch.tensor([], dtype=torch.int64), dist_reduce_fx="cat")
-        self.add_state("target", default=torch.tensor([], dtype=torch.int64), dist_reduce_fx="cat")
+        self.add_state("idx", default=[], dist_reduce_fx=None)
+        self.add_state("preds", default=[], dist_reduce_fx=None)
+        self.add_state("target", default=[], dist_reduce_fx=None)
 
-    def update(self, idx: torch.Tensor, preds: torch.Tensor, target: torch.Tensor):
-        assert idx.shape == preds.shape == target.shape, (
-            f"Indexes, predicions and targets must be of the same shape"
-        )
+    def update(self, idx: torch.Tensor, preds: torch.Tensor, target: torch.Tensor) -> None:
+        if not (idx.shape == target.shape == preds.shape):
+            raise ValueError("`idx`, `preds` and `target` must be of the same shape")
+
+        idx = idx.to(dtype=torch.int64).flatten()
+        preds = preds.to(dtype=torch.float32).flatten()
+        target = target.to(dtype=torch.int64).flatten()
+
+        self.idx.append(idx)
+        self.preds.append(preds)
+        self.target.append(target)
+
+    def compute(self) -> torch.Tensor:
+        r"""
+        First concat state `idx`, `preds` and `target` since they were stored as lists. After that,
+        compute list of groups that will help in keeping together predictions about the same query.
+        Finally, for each group compute the `_metric` if the number of positive targets is at least
+        1, otherwise behave as specified by `self.query_without_relevant_docs`.
+        """
+
+        idx = torch.cat(self.idx, dim=0)
+        preds = torch.cat(self.preds, dim=0)
+        target = torch.cat(self.target, dim=0)
 
-        self.idx = torch.cat([self.idx, idx])
-        self.preds = torch.cat([self.preds, preds])
-        self.target = torch.cat([self.target, target])
-    
-    def compute(self):
         res = []
-        for group in get_mini_groups(self.idx):
-            if self.target[group].sum() == 0:
-                if self.empty_documents == 'error':
+        kwargs = {'device': idx.device, 'dtype': torch.float32}
+
+        groups = get_mini_groups(idx)
+        for group in groups:
+
+            mini_preds = preds[group]
+            mini_target = target[group]
+
+            if not mini_target.sum():
+                if self.query_without_relevant_docs == 'error':
                     raise ValueError(
-                        f"{self.__class__.__name__} was provided with a prediction with no positive values, idx: {group}"
+                        f"`{self.__class__.__name__}.compute()` was provided with "
+                        f"a query without positive targets, indexes: {group}"
                     )
-                elif self.empty_documents == 'positive': res.append(torch.tensor(1.0))
-                elif self.empty_documents == 'negative': res.append(torch.tensor(0.0))
+                if self.query_without_relevant_docs == 'pos':
+                    res.append(torch.tensor(1.0, **kwargs))
+                elif self.query_without_relevant_docs == 'neg':
+                    res.append(torch.tensor(0.0, **kwargs))
             else:
-                res.append(
-                    self.metric(group)
-                )
-        return torch.stack(res).mean()
-
-    def metric(self, group: List[int]):
-        r""" Compute a metric over a single group. """
-        raise NotImplementedError("This method must be overridden by subclasses")
+                res.append(self._metric(mini_preds, mini_target))
+
+        if len(res) > 0:
+            return torch.stack(res).mean()
+        return torch.tensor(0.0, **kwargs)
+
+    @abstractmethod
+    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        r"""
+        Compute a metric over a predictions and target of a single group.
+        This method should be overridden by subclasses.
+        """
diff --git a/transformers_lightning/models/__init__.py b/transformers_lightning/models/__init__.py
index 7dcf38d..9e3b002 100644
--- a/transformers_lightning/models/__init__.py
+++ b/transformers_lightning/models/__init__.py
@@ -1,1 +1,1 @@
-from transformers_lightning.models.transformers_model import TransformersModel
\ No newline at end of file
+from transformers_lightning.models.transformers_model import TransformersModel
diff --git a/transformers_lightning/models/transformers_model.py b/transformers_lightning/models/transformers_model.py
index d0afcbf..2df0a6b 100644
--- a/transformers_lightning/models/transformers_model.py
+++ b/transformers_lightning/models/transformers_model.py
@@ -55,24 +55,29 @@ def configure_optimizers(self):
         ]
 
         # Define adam optimizer
-        optimizer = AdamW(optimizer_grouped_parameters,
-                          lr=self.hparams.learning_rate,
-                          eps=self.hparams.adam_epsilon,
-                          betas=self.hparams.adam_betas)
+        optimizer = AdamW(
+            optimizer_grouped_parameters,
+            lr=self.hparams.learning_rate,
+            eps=self.hparams.adam_epsilon,
+            betas=self.hparams.adam_betas
+        )
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = LinearSchedulerWithWarmup(optimizer,
-                                              num_warmup_steps=self.hparams.warmup_steps,
-                                              num_training_steps=max_steps,
-                                              beg_step=self.hparams.beg_scheduler_step)
+        scheduler = LinearSchedulerWithWarmup(
+            optimizer,
+            num_warmup_steps=self.hparams.warmup_steps,
+            num_training_steps=max_steps,
+            beg_step=self.hparams.beg_scheduler_step
+        )
 
         return {
             'optimizer': optimizer,
-            'lr_scheduler': {
-                'scheduler': scheduler, # The LR schduler
-                'interval': 'step', # The unit of the scheduler's step size
-                'frequency': 1, # The frequency of the scheduler
-            }
+            'lr_scheduler':
+                {
+                    'scheduler': scheduler,    # The LR schduler
+                    'interval': 'step',    # The unit of the scheduler's step size
+                    'frequency': 1,    # The frequency of the scheduler
+                }
         }
 
     @staticmethod
diff --git a/transformers_lightning/optimizers/README.md b/transformers_lightning/optimizers/README.md
new file mode 100644
index 0000000..6d8427a
--- /dev/null
+++ b/transformers_lightning/optimizers/README.md
@@ -0,0 +1,8 @@
+# Optimizers
+
+Optimizers manage weight update starting from gradient values. They may have complex internal states to better move on the loss multi-dimensional surface.
+
+
+## ElectraAdamW
+
+This optimizer is same as `AdamW` but for a small fix to the moving average update. Original implementation can be found [here](https://github.com/google-research/electra/blob/f93f3f81cdc13435dd3e85766852d00ff3e00ab5/model/optimization.py#L70).
diff --git a/transformers_lightning/optimizers/__init__.py b/transformers_lightning/optimizers/__init__.py
new file mode 100644
index 0000000..75315b1
--- /dev/null
+++ b/transformers_lightning/optimizers/__init__.py
@@ -0,0 +1,1 @@
+from transformers_lightning.optimizers.adamw_electra import ElectraAdamW
diff --git a/transformers_lightning/optimizers/adamw_electra.py b/transformers_lightning/optimizers/adamw_electra.py
new file mode 100644
index 0000000..d105c40
--- /dev/null
+++ b/transformers_lightning/optimizers/adamw_electra.py
@@ -0,0 +1,103 @@
+import math
+
+import torch
+from torch.optim import AdamW
+
+
+class ElectraAdamW(AdamW):
+    r"""Implements AdamW algorithm.
+
+    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay coefficient (default: 1e-2)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-4, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, amsgrad=False):
+        super(ElectraAdamW,
+              self).__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                # Perform stepweight decay
+                p.mul_(1 - group['lr'] * group['weight_decay'])
+
+                # Perform optimization step
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError('AdamW does not support sparse gradients')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+                bias_correction1 = 1 - beta1**state['step']
+                bias_correction2 = 1 - beta2**state['step']
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).add_(grad.sqrt(), value=1 - beta2)    # <-- fix is here
+
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
+                else:
+                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
+
+                step_size = group['lr'] / bias_correction1
+
+                p.addcdiv_(exp_avg, denom, value=-step_size)
+
+        return loss
diff --git a/transformers_lightning/schedulers/README.md b/transformers_lightning/schedulers/README.md
index 9a32635..da1ae5b 100644
--- a/transformers_lightning/schedulers/README.md
+++ b/transformers_lightning/schedulers/README.md
@@ -1,6 +1,6 @@
 # Schedulers
 
-Schedulers modify the learning rate through the epochs.
+Schedulers modify the learning rate through the epochs. *Every* scheduler is designed to be called every *step* and *not* every *epoch*.
 
 
 ## LinearScheduler
@@ -31,3 +31,10 @@ This scheduler will linearly increment the learning rate from 0 to the `learning
 ## CosineSchedulerWithWarmupAndHardRestart
 
 This scheduler will linearly increment the learning rate from 0 to the `learning_rate` in the first `num_warmup_steps` steps. After that, it will decrease it up to `0` following the value of the cosine function with several hard restarts.
+
+
+## PolynomialLayerwiseDecaySchedulerWithWarmup
+
+Create a polynomially decreasing scheduler. Conversion of `https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/polynomial_decay`. More informations about the default parameters can be found on the documentation of `_LRScheduler` in the `torch` project.
+
+The `layerwise_lr_decay_power` parameter allows to multiply the learning rate for each group by a factor `f = layerwise_lr_decay_power^(max(depth) - depth)`. `depth` is a key that should be defined in every group of parameters, along with the usual `weight_decay`. A small `depth` will mean that the learning rate will be only slightly decreased. A high `depth` will reduce a lot the effects of the training on that parameters.
diff --git a/transformers_lightning/schedulers/__init__.py b/transformers_lightning/schedulers/__init__.py
index a3683a5..a50b91a 100644
--- a/transformers_lightning/schedulers/__init__.py
+++ b/transformers_lightning/schedulers/__init__.py
@@ -3,4 +3,5 @@
 from transformers_lightning.schedulers.constant_scheduler import ConstantScheduler
 from transformers_lightning.schedulers.constant_scheduler_with_warmup import ConstantSchedulerWithWarmup
 from transformers_lightning.schedulers.cosine_scheduler_with_warmup import CosineSchedulerWithWarmup
-from transformers_lightning.schedulers.cosine_scheduler_with_warmup_and_hard_restart import CosineSchedulerWithWarmupAndHardRestart
\ No newline at end of file
+from transformers_lightning.schedulers.cosine_scheduler_with_warmup_and_hard_restart import CosineSchedulerWithWarmupAndHardRestart
+from transformers_lightning.schedulers.layerwise_decay_scheduler import PolynomialLayerwiseDecaySchedulerWithWarmup
diff --git a/transformers_lightning/schedulers/constant_scheduler.py b/transformers_lightning/schedulers/constant_scheduler.py
index 7804ea0..e45e659 100644
--- a/transformers_lightning/schedulers/constant_scheduler.py
+++ b/transformers_lightning/schedulers/constant_scheduler.py
@@ -19,17 +19,11 @@ class ConstantScheduler(_LRScheduler):
         >>> scheduler = ConstantScheduler(optimizer)
     """
 
-    def __init__(
-        self,
-        optimizer: torch.optim.Optimizer,
-        last_epoch: int = -1,
-        verbose: bool = False
-    ):
+    def __init__(self, optimizer: torch.optim.Optimizer, last_epoch: int = -1, verbose: bool = False):
         super().__init__(optimizer, last_epoch, verbose)
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
         return [base_lr for base_lr in self.base_lrs]
diff --git a/transformers_lightning/schedulers/constant_scheduler_with_warmup.py b/transformers_lightning/schedulers/constant_scheduler_with_warmup.py
index a521c44..edcc7f9 100644
--- a/transformers_lightning/schedulers/constant_scheduler_with_warmup.py
+++ b/transformers_lightning/schedulers/constant_scheduler_with_warmup.py
@@ -22,15 +22,11 @@ class ConstantSchedulerWithWarmup(_LRScheduler):
     """
 
     def __init__(
-        self,
-        optimizer: torch.optim.Optimizer,
-        num_warmup_steps: int = 0,
-        last_epoch: int = -1,
-        verbose: bool = False
+        self, optimizer: torch.optim.Optimizer, num_warmup_steps: int = 0, last_epoch: int = -1, verbose: bool = False
     ):
         if not isinstance(num_warmup_steps, int) or not num_warmup_steps >= 0:
             raise ValueError("`num_warmup_steps` must be an integer greater than 0")
-        
+
         self.num_warmup_steps = num_warmup_steps
 
         super().__init__(optimizer, last_epoch, verbose)
@@ -42,8 +38,6 @@ def lr_lambda(self, current_step: int) -> int:
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
-        return [base_lr * self.lr_lambda(self.last_epoch)
-                for base_lr in self.base_lrs]
\ No newline at end of file
+        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]
diff --git a/transformers_lightning/schedulers/cosine_scheduler_with_warmup.py b/transformers_lightning/schedulers/cosine_scheduler_with_warmup.py
index 5a2fc9d..0d804d0 100644
--- a/transformers_lightning/schedulers/cosine_scheduler_with_warmup.py
+++ b/transformers_lightning/schedulers/cosine_scheduler_with_warmup.py
@@ -51,13 +51,12 @@ def __init__(
     def lr_lambda(self, current_step):
         if current_step < self.num_warmup_steps:
             return float(current_step) / float(max(1, self.num_warmup_steps))
-        progress = float(current_step - self.num_warmup_steps) / float(max(1, self.num_training_steps - self.num_warmup_steps))
+        progress = float(current_step -
+                         self.num_warmup_steps) / float(max(1, self.num_training_steps - self.num_warmup_steps))
         return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)))
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
-        return [base_lr * self.lr_lambda(self.last_epoch)
-                for base_lr in self.base_lrs]
\ No newline at end of file
+        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]
diff --git a/transformers_lightning/schedulers/cosine_scheduler_with_warmup_and_hard_restart.py b/transformers_lightning/schedulers/cosine_scheduler_with_warmup_and_hard_restart.py
index 77d47da..e172ee7 100644
--- a/transformers_lightning/schedulers/cosine_scheduler_with_warmup_and_hard_restart.py
+++ b/transformers_lightning/schedulers/cosine_scheduler_with_warmup_and_hard_restart.py
@@ -51,15 +51,14 @@ def __init__(
     def lr_lambda(self, current_step):
         if current_step < self.num_warmup_steps:
             return float(current_step) / float(max(1, self.num_warmup_steps))
-        progress = float(current_step - self.num_warmup_steps) / float(max(1, self.num_training_steps - self.num_warmup_steps))
+        progress = float(current_step -
+                         self.num_warmup_steps) / float(max(1, self.num_training_steps - self.num_warmup_steps))
         if progress >= 1.0:
             return 0.0
         return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(self.num_cycles) * progress) % 1.0))))
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
-        return [base_lr * self.lr_lambda(self.last_epoch)
-                for base_lr in self.base_lrs]
+        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]
diff --git a/transformers_lightning/schedulers/layerwise_decay_scheduler.py b/transformers_lightning/schedulers/layerwise_decay_scheduler.py
new file mode 100644
index 0000000..a6c0e94
--- /dev/null
+++ b/transformers_lightning/schedulers/layerwise_decay_scheduler.py
@@ -0,0 +1,131 @@
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+# Converted to PyTorch
+"""Functions and classes related to optimization (weight updates).
+Modified from the original BERT code to allow for having separate learning
+rates for different layers of the network.
+"""
+import math
+
+import warnings
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class PolynomialLayerwiseDecaySchedulerWithWarmup(_LRScheduler):
+    r"""
+    Create a polynomially decreasing scheduler.
+    Conversion of `https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/polynomial_decay`.
+    More informations about the default parameters can be found on the documentation of `_LRScheduler` in the `torch` project.
+
+    If `layerwise_lr_decay_power` is different from 1.0, the learning rate over each group will
+    be multiplied by a factor `f = layerwise_lr_decay_power^(max(depth) - depth)`. `depth` is a
+    key that should be defined in every group of parameters, along with the usual `weight_decay`.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_training_steps (:obj:`int`):
+            number of steps in which lr is decayed.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+        verbose (:obj:`bool`, `optional`, defaults to False): 
+            If ``True``, prints a message to stdout for each update.
+        end_learning_rate (:obj:`float`, `optional`, defaults to 0.0001):
+            target learning rate in the last step/epoch.
+        lr_decay_power (:obj:`float`, `optional`, defaults to 1.0):
+            learning rate decay base for polynomial decay.
+        layerwise_lr_decay_power (:obj:`float`, `optional`, defaults to 1.0):
+            learning rate decay base for layerwise decay. 
+        cycle (:obj:`bool`, `optional`, defaults to False):
+            whether to extend decay steps when global step is higher.
+        warmup_steps (:obj:`int`, `optional`, defaults to 0):
+            number of steps of the warmup phase.
+
+    Example:
+        >>> scheduler = PolynomialLayerwiseDecaySchedulerWithWarmup(optimizer)
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        num_training_steps: int,
+        last_epoch: int = -1,
+        verbose: bool = False,
+        end_learning_rate: float = 0.0001,
+        lr_decay_power: float = 1.0,
+        layerwise_lr_decay_power: float = 1.0,
+        cycle: bool = False,
+        warmup_steps: int = 0
+    ):
+        if num_training_steps < 0:
+            raise ValueError("`num_training_steps` must be an integer greater than 0")
+
+        if end_learning_rate < 0:
+            raise ValueError(f"Cannot define negative `end_learning_rate`, found {end_learning_rate}")
+
+        if layerwise_lr_decay_power < 0:
+            raise ValueError(f"Cannot define negative `layerwise_lr_decay_power`, found {layerwise_lr_decay_power}")
+
+        if warmup_steps < 0:
+            raise ValueError(f"Cannot define negative `warmup_steps`, found {warmup_steps}")
+
+        self.num_training_steps = num_training_steps
+        self.end_learning_rate = end_learning_rate
+        self.lr_decay_power = lr_decay_power
+        self.cycle = cycle
+        self.warmup_steps = warmup_steps
+        self.layerwise_lr_decay_power = layerwise_lr_decay_power
+
+        # retrieve depth for each params group
+        self.depths = [group['depth'] for group in self.optimizer.param_groups]
+
+        super().__init__(optimizer, last_epoch=last_epoch, verbose=verbose)
+
+    def _layerwise_decay(self, lrs):
+        """
+        Have lower learning rates for layers closer to the input.
+        Requires that groups passed to the Optimizer are already sorted from the
+        closer to the input to the closer the output.
+        """
+
+        return [lr * (self.layerwise_lr_decay_power**(max(self.depths) - depth)) for lr, depth in zip(lrs, self.depths)]
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`.", UserWarning
+            )
+
+        decay_steps = self.num_training_steps
+        global_step = self.last_epoch    # often last_epoch is used to indicate the number of steps. It is updated every `sched.step()`...
+
+        # if cycle, extend decay_steps if larger than global_step
+        if self.cycle:
+            decay_steps = decay_steps * math.ceil(global_step / decay_steps)
+        else:
+            global_step = min(global_step, decay_steps)
+
+        lrs = [
+            (
+                (base_lr - self.end_learning_rate) *
+                (1 - global_step / decay_steps)**(self.lr_decay_power) + self.end_learning_rate
+            ) * (min(1.0, global_step / self.warmup_steps) if self.warmup_steps > 0 else 1.0)
+            for base_lr in self.base_lrs
+        ]
+
+        lrs = self._layerwise_decay(lrs)
+        return lrs
diff --git a/transformers_lightning/schedulers/linear_scheduler.py b/transformers_lightning/schedulers/linear_scheduler.py
index 30107e1..1249834 100644
--- a/transformers_lightning/schedulers/linear_scheduler.py
+++ b/transformers_lightning/schedulers/linear_scheduler.py
@@ -22,28 +22,20 @@ class LinearScheduler(_LRScheduler):
     """
 
     def __init__(
-        self,
-        optimizer: torch.optim.Optimizer,
-        num_training_steps: int,
-        last_epoch: int = -1,
-        verbose: bool = False
+        self, optimizer: torch.optim.Optimizer, num_training_steps: int, last_epoch: int = -1, verbose: bool = False
     ):
         if not isinstance(num_training_steps, int) or not num_training_steps >= 0:
             raise ValueError("`num_training_steps` must be an integer greater than 0")
-        
+
         self.num_training_steps = num_training_steps
 
         super().__init__(optimizer, last_epoch, verbose)
 
     def lr_lambda(self, current_step: int) -> int:
-        return max(
-            0.0, float(self.num_training_steps - current_step) / self.num_training_steps
-        )
+        return max(0.0, float(self.num_training_steps - current_step) / self.num_training_steps)
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
-        return [base_lr * self.lr_lambda(self.last_epoch)
-                for base_lr in self.base_lrs]
+        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]
diff --git a/transformers_lightning/schedulers/linear_scheduler_with_warmup.py b/transformers_lightning/schedulers/linear_scheduler_with_warmup.py
index bd7b456..2cb88d1 100644
--- a/transformers_lightning/schedulers/linear_scheduler_with_warmup.py
+++ b/transformers_lightning/schedulers/linear_scheduler_with_warmup.py
@@ -40,13 +40,12 @@ def __init__(
 
         if not isinstance(num_warmup_steps, int) or not num_warmup_steps >= 0:
             raise ValueError("`num_warmup_steps` must be an integer greater than 0")
-        
+
         self._num_warmup_steps = num_warmup_steps
         self._num_training_steps = num_training_steps
         self._beg_step = beg_step
 
-
-        super().__init__(optimizer, last_epoch, verbose)
+        super().__init__(optimizer, last_epoch=last_epoch, verbose=verbose)
 
     def lr_lambda(self, current_step: int) -> int:
 
@@ -59,13 +58,13 @@ def lr_lambda(self, current_step: int) -> int:
         if relative_curr_step < self._num_warmup_steps:
             return float(relative_curr_step) / float(max(1, self._num_warmup_steps))
         return max(
-            0.0, float(relative_num_training_steps - relative_curr_step) / float(max(1, relative_num_training_steps - self._num_warmup_steps))
+            0.0,
+            float(relative_num_training_steps - relative_curr_step) /
+            float(max(1, relative_num_training_steps - self._num_warmup_steps))
         )
 
     def get_lr(self):
         if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
+            warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.")
 
-        return [base_lr * self.lr_lambda(self.last_epoch)
-                for base_lr in self.base_lrs]
+        return [base_lr * self.lr_lambda(self.last_epoch) for base_lr in self.base_lrs]
diff --git a/transformers_lightning/utils/distributed.py b/transformers_lightning/utils/distributed.py
index eae1812..f97db8c 100644
--- a/transformers_lightning/utils/distributed.py
+++ b/transformers_lightning/utils/distributed.py
@@ -44,12 +44,9 @@ def compute_max_steps(hparams: Namespace, trainer: Trainer) -> int:
     # if cannot retrieve len of the dataset, skip
     # this can happen with iterabledatasets
     if not (
-        hasattr(trainer.datamodule.train_dataset, '__len__') or
-        hasattr(trainer.datamodule.train_dataset, 'length')
+        hasattr(trainer.datamodule.train_dataset, '__len__') or hasattr(trainer.datamodule.train_dataset, 'length')
     ):
-        logger.warning(
-            "Cannot infer dataset length from datamodule.train_dataset, returning max_steps=None"
-        )
+        logger.warning("Cannot infer dataset length from datamodule.train_dataset, returning max_steps=None")
         return None
 
     try:
@@ -67,8 +64,6 @@ def compute_max_steps(hparams: Namespace, trainer: Trainer) -> int:
     steps_per_epoch = math.ceil(training_batches_per_epoch / hparams.accumulate_grad_batches)
     steps = hparams.max_epochs * steps_per_epoch
 
-    logger.warning(
-        f"Automagically computed max_steps={steps}. If it appears to be OK, ignore this warning"
-    )
+    logger.warning(f"Automagically computed max_steps={steps}. If it appears to be OK, ignore this warning")
 
     return steps
diff --git a/transformers_lightning/utils/functional.py b/transformers_lightning/utils/functional.py
index 0dd4cfd..8cc55af 100644
--- a/transformers_lightning/utils/functional.py
+++ b/transformers_lightning/utils/functional.py
@@ -12,6 +12,7 @@ def remove_from_dict(dictionary: Dict, keys: List[str] = []) -> Dict:
         assert k in dictionary, f"{k} not in dict {dictionary}"
         del dictionary[k]
 
+
 def all_equal_in_iterable(iterable: Iterable) -> bool:
     r"""
     Check all items in an iterable 
@@ -21,12 +22,14 @@ def all_equal_in_iterable(iterable: Iterable) -> bool:
     else:
         return iterable.count(iterable[0]) == len(iterable)
 
+
 def flatten(list_of_lists: List[List]) -> List:
     r"""
     Unroll a list of lists
     """
     return [item for sublist in list_of_lists for item in sublist]
 
+
 def get_inner_type(parameter_list: List) -> type:
     r"""
     Check both that all entries in a list are of the same and return that type
@@ -39,6 +42,7 @@ def get_inner_type(parameter_list: List) -> type:
             assert parameter_type is type(param)
     return parameter_type
 
+
 def model_checksum(model: torch.nn.Module) -> torch.Tensor:
     r"""
     Return a checksum of a torch model by summing all the paramter values.
@@ -47,6 +51,7 @@ def model_checksum(model: torch.nn.Module) -> torch.Tensor:
     """
     return sum([x.data.sum() for x in model.parameters()])
 
+
 def model_gradient_checksum(model: torch.nn.Module) -> torch.Tensor:
     r"""
     Return a checksum of the gradients of a torch model by summing all the gradient values.
@@ -55,6 +60,7 @@ def model_gradient_checksum(model: torch.nn.Module) -> torch.Tensor:
     """
     return sum([x.grad.sum() if x.grad is not None else 0.0 for x in model.parameters()])
 
+
 def concat_dict_values(data: List) -> Dict:
     r"""
     Given a list of dictionaries with the same keys, return a dictionary in which
@@ -69,6 +75,7 @@ def concat_dict_values(data: List) -> Dict:
                 res[key] = [dictionary[key]]
     return res
 
+
 def split(a: List, n: int) -> tuple:
     r"""
     Split a list in `n` equal parts (or similar, if `len(a) % n != 0`)
@@ -76,6 +83,7 @@ def split(a: List, n: int) -> tuple:
     k, m = divmod(len(a), n)
     return (a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))
 
+
 def safe_merge(dict_1: Dict, dict_2: Dict) -> Dict:
     r"""
     Merge two dictionaries asserting that they do not have common keys
@@ -89,6 +97,7 @@ def safe_merge(dict_1: Dict, dict_2: Dict) -> Dict:
         res[key] = value
     return res
 
+
 def collate_single_fn(data: List[Dict]) -> Dict[str, torch.Tensor]:
     r"""
     Merge n dicts with identical keys creating list of value tensors.
@@ -98,6 +107,7 @@ def collate_single_fn(data: List[Dict]) -> Dict[str, torch.Tensor]:
     res = {k: torch.tensor(v) for k, v in res.items()}
     return res
 
+
 def collate_multi_fn(data):
     r"""
     Concatenate the values of each batch dict of each dataset when using, for example, a `StackDataset`
@@ -105,6 +115,7 @@ def collate_multi_fn(data):
     res = [collate_single_fn(x) for x in zip(*data)]
     return res
 
+
 def concat_generators(*args: Iterable[Generator]) -> Generator:
     r"""
     Concat generators by yielding from first, second, ..., n-th
@@ -112,6 +123,7 @@ def concat_generators(*args: Iterable[Generator]) -> Generator:
     for gen in args:
         yield from gen
 
+
 def join_on_path(files_list: List[str], base_path: str) -> List[str]:
     r"""
     Join a list of files on a given path
@@ -121,22 +133,24 @@ def join_on_path(files_list: List[str], base_path: str) -> List[str]:
             type(files_list), type(base_path))
     return [os.path.join(base_path, f) for f in files_list]
 
+
 def uniform_distrib(size: Union[Tuple, List, torch.Size], **kwargs) -> torch.Tensor:
     r"""
     Return a vector containing uniform values such that the last axis sums to 1.0.
     """
-    return torch.full(size=size, fill_value=(1/size[-1]), **kwargs)
+    return torch.full(size=size, fill_value=(1 / size[-1]), **kwargs)
+
 
 def normalize_standard(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
     r"""
     Do (x - E[x]) / ( E[x - E[x]] ) on axis specified by dim.
     """
     return torch.where(
-        (x.std(dim=dim) > 0).unsqueeze(-1),
-        (x - x.mean(dim=dim).unsqueeze(-1)) / x.std(dim=dim).unsqueeze(-1),
+        (x.std(dim=dim) > 0).unsqueeze(-1), (x - x.mean(dim=dim).unsqueeze(-1)) / x.std(dim=dim).unsqueeze(-1),
         (x - x.mean(dim=dim).unsqueeze(-1))
     )
 
+
 def normalize_linear(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
     r"""
     Do (x - min(x)) / (max(x) - min(x)) on axis specified by dim.
@@ -145,6 +159,7 @@ def normalize_linear(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
     min_x = torch.min(x, dim=dim)
     return (x - min_x) / (max_x - min_x)
 
+
 def concat_dict_tensors(*args: Iterable[Dict], dim: int = 0) -> Dict[str, torch.Tensor]:
     r"""
     Concat dictionaries values containing tensors on `dim` dimension.
@@ -162,4 +177,3 @@ def concat_dict_tensors(*args: Iterable[Dict], dim: int = 0) -> Dict[str, torch.
 
     res = {k: torch.cat(v, dim=dim) for k, v in res.items()}
     return res
-
diff --git a/transformers_lightning/utils/generators.py b/transformers_lightning/utils/generators.py
index d65c2bd..8f210e2 100644
--- a/transformers_lightning/utils/generators.py
+++ b/transformers_lightning/utils/generators.py
@@ -33,6 +33,7 @@ def filter_generator(generator_in: Generator, step: int = 1, offset: int = 0) ->
         except:
             return
 
+
 def batch_filter(generator_in: Generator, size: int = 1) -> Generator:
     r"""
     By reading `size` elements at a time, we assure that no last iteration will have
diff --git a/transformers_lightning/utils/gradient.py b/transformers_lightning/utils/gradient.py
index 4522192..649477e 100644
--- a/transformers_lightning/utils/gradient.py
+++ b/transformers_lightning/utils/gradient.py
@@ -7,6 +7,7 @@ class RevGrad(Function):
 
     Reverse the gradient values (* -1) while doing backpropagation.
     """
+
     @staticmethod
     def forward(ctx, input_):
         ctx.save_for_backward(input_)
@@ -14,13 +15,11 @@ def forward(ctx, input_):
         return output
 
     @staticmethod
-    def backward(ctx, grad_output):  # pragma: no cover
+    def backward(ctx, grad_output):    # pragma: no cover
         grad_input = None
         if ctx.needs_input_grad[0]:
             grad_input = -grad_output
         return grad_input
 
-revgrad = RevGrad.apply
-
-
 
+revgrad = RevGrad.apply
diff --git a/transformers_lightning/utils/inspectors.py b/transformers_lightning/utils/inspectors.py
index 3873f73..2363508 100644
--- a/transformers_lightning/utils/inspectors.py
+++ b/transformers_lightning/utils/inspectors.py
@@ -11,6 +11,7 @@ def get_classes_from_module(config_proto, parent=None, do_lower_case=False):
                 res[i[0].lower()] = i[1]
     return res
 
+
 def get_classes_from_module_list(configs_proto, parent=None):
     res = {}
     for config_proto in configs_proto:
@@ -20,6 +21,7 @@ def get_classes_from_module_list(configs_proto, parent=None):
                 res[i[0]] = i[1]
     return res
 
+
 def get_functions_from_module(config_proto):
     res = {}
     for i in getmembers(config_proto):
@@ -27,6 +29,7 @@ def get_functions_from_module(config_proto):
             res[i[0]] = i[1]
     return res
 
+
 def get_types_from_module(config_proto, target_type=tuple):
     res = {}
     variables = [x for x in dir(config_proto) if not x.startswith("_")]
@@ -35,5 +38,6 @@ def get_types_from_module(config_proto, target_type=tuple):
             res[i[0]] = i[1]
     return res
 
+
 def is_simple(thing):
-    return not isfunction(thing) and not isclass(thing)
\ No newline at end of file
+    return not isfunction(thing) and not isclass(thing)
diff --git a/transformers_lightning/utils/language_modeling.py b/transformers_lightning/utils/language_modeling.py
index f1f68aa..1209cb2 100644
--- a/transformers_lightning/utils/language_modeling.py
+++ b/transformers_lightning/utils/language_modeling.py
@@ -1,176 +1,4 @@
-import numpy as np
 import torch
-import transformers
-import scipy.stats as st
-from transformers_lightning import utils
-from pytorch_lightning import _logger as logger
-
-IGNORE_IDX = -100
-
-
-def mask_tokens_for_mlm(
-    inputs: torch.Tensor,
-    tokenizer: transformers.PreTrainedTokenizer,
-    mlm_probability: float = 0.15,
-    weights: torch.Tensor = None,
-    reliability: float = 0.05,
-    **kwargs
-) -> tuple:
-    """
-    Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
-    If weights are provided, probability of each token to be masked will be weighted.
-    `inputs` and `weights` must be on the same device or `auto_move_tensors` must be true!
-    """
-
-    if 'auto_move_tensors' in kwargs:
-        logger.warning("`auto_move_tensors` is deprecated and will be removed in 0.3.0")
-
-    if tokenizer.mask_token is None:
-        raise ValueError(
-            "This tokenizer does not have a mask token which is necessary for masked language modeling. "
-            "Remove the --mlm flag if you want to use this tokenizer."
-        )
-
-    device = inputs.device
-
-    labels = inputs.clone()
-    # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
-    if weights is None:
-        probability_matrix = torch.full(labels.shape, fill_value=mlm_probability, dtype=torch.float32, device=device)
-    else:
-        weights = weights.to(device=labels.device)
-
-        """
-        Going to use the weights in the following way:
-        - w are the weights
-        - p the probability of masking a token
-
-        probabilities = ( 1 + (w - w_mean) / (w_std * Z_(1 - r/2)) ) -> clipped in [0, 1]
-        """
-        z_score = st.norm.ppf(1 - reliability/2)
-        probability_matrix = mlm_probability * (1 + utils.normalize_standard(weights[labels], dim=-1) / z_score)
-        probability_matrix = torch.clip(probability_matrix, min=0, max=1)
-
-    special_tokens_mask = [
-        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
-    ]
-    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool, device=device), value=0.0)
-
-    if tokenizer._pad_token is not None:
-        padding_mask = labels.eq(tokenizer.pad_token_id)
-        probability_matrix.masked_fill_(padding_mask, value=0.0)
-    masked_indices = torch.bernoulli(probability_matrix).bool()
-    labels[~masked_indices] = IGNORE_IDX  # We only compute loss on masked tokens
-
-    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
-    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8, device=device)).bool() & masked_indices
-    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
-
-    # 10% of the time, we replace masked input tokens with random word
-    indices_random = torch.bernoulli(
-        torch.full(labels.shape, 0.5, device=device)).bool() & masked_indices & ~indices_replaced
-    random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long, device=device)
-    inputs[indices_random] = random_words[indices_random]
-
-    # The rest of the time (10% of the time) we keep the masked input tokens unchanged
-    return inputs, labels
-
-
-def random_token_substutition(
-    inputs: torch.Tensor,
-    tokenizer: transformers.PreTrainedTokenizer,
-    rts_probability: float = 0.15,
-    do_not_touch_input_ids=False,
-    weights: torch.Tensor = None,
-    reliability: float = 0.05,
-    ngrams:int = 1,
-    **kwargs
-) -> tuple:
-    """
-    Prepare tokens inputs/labels for random token substutition modeling.
-    We sample a few tokens in each sequence for RTS training
-    (with probability args.rts_probability defaults to 0.15 in Bert/RoBERTa)
-    
-    Example
-    >>> from transformers import BertTokenizer
-    >>> tok = BertTokenizer.from_pretrained("bert-base-cased")
-
-    >>> bs = tok.encode_plus("ciao sono luca")
-    >>> bs = {k: torch.tensor(v).unsqueeze(0) for k,v in bs.items()}
-    >>> print(bs)
-    {
-        'input_ids': tensor([[  101,   172,  1465,  1186,  1488,  1186,   181, 23315,   102]]),
-        'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0]]),
-        'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1]])
-    }
-
-    >>> bs["input_ids"], bs["labels"] = random_token_substutition(
-    >>>     bs["input_ids"], tok, rts_probability=0.4
-    >>> )
-    >>> print(bs)
-    {
-        'input_ids': tensor([[  101,   172,  8959,  1186,  3082, 27735,   181,  1348,   102]]),
-        'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0]]),
-        'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1]]),
-        'labels': tensor([[IGNORE_IDX, 0, 1, 0, 1, 1, 0, 1, IGNORE_IDX]])
-    }
-
-    If `do_not_touch_input_ids` is True, input ids are not modified and only labels are created.
-    If `weights` if not None, the probability matrix will be initialized giving each token a probability
-    proportional to the weight. `weights` should be a vector on length `vocab_size`.
-    """
-    if 'auto_move_tensors' in kwargs:
-        logger.warning("`auto_move_tensors` is deprecated and will be removed in 0.3.0")
-
-    # creating tensor directly on right device is more efficient
-    device = inputs.device
-
-    # We sample a few tokens in each sequence for masked-LM training (with probability args.rts_probability defaults to 0.15 in Bert/RoBERTa)
-    if weights is None:
-        probability_matrix = torch.full(inputs.shape, fill_value=rts_probability, dtype=torch.float32, device=device)
-    else:
-        weights = weights.to(device=device)
-
-        """
-        Going to use the weights in the following way:
-        - w are the weights
-        - p the probability of masking a token
-
-        probabilities = ( 1 + (w - w_mean) / (w_std * Z_(1 - r/2)) ) -> clipped in [0, 1]
-        """
-        z_score = st.norm.ppf(1 - reliability/2)
-        probability_matrix = rts_probability * (1 + utils.normalize_standard(weights[inputs], dim=-1) / z_score)
-        probability_matrix = torch.clip(probability_matrix, min=0, max=1)
-
-    # probability matrix containts the probability of each token to be randomly substituted
-    labels = torch.full(inputs.shape, fill_value=0, dtype=torch.long, device=device)
-
-    # not going to substitute special tokens of the LM (bert, roby, ...)
-    special_tokens_mask = [
-        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in inputs.tolist()
-    ]
-    special_tokens_mask_tensor = torch.tensor(special_tokens_mask, dtype=torch.bool, device=device)
-    probability_matrix.masked_fill_(special_tokens_mask_tensor, value=0.0)
-    labels.masked_fill_(special_tokens_mask_tensor, value=IGNORE_IDX)
-
-    # no need to substitute padding tokens, assigning 0.0 prob
-    if tokenizer._pad_token is not None:
-        padding_mask = inputs.eq(tokenizer.pad_token_id)
-        probability_matrix.masked_fill_(padding_mask, value=0.0)
-        labels.masked_fill_(padding_mask, value=IGNORE_IDX)
-
-    substituted_indices = torch.bernoulli(probability_matrix).bool()
-
-    random_words = torch.randint(len(tokenizer), inputs.shape, dtype=torch.long, device=device)
-    if not do_not_touch_input_ids:
-        inputs[substituted_indices] = random_words[substituted_indices]
-    labels.masked_fill_(substituted_indices, value=1)
-
-    # compute the ngrams labels
-    if ngrams > 1:
-        labels = make_ngrams_labels(labels, ngrams, ignore_idx=IGNORE_IDX)
-
-    return inputs, labels
 
 
 def make_ngrams_labels(labels, n, ignore_idx=None):
@@ -178,19 +6,19 @@ def make_ngrams_labels(labels, n, ignore_idx=None):
     if ignore_idx is None:
         grams = labels.repeat(1, n).view(labels.size(0), n, -1)
     else:
-        grams = labels[:,1:].repeat(1, n).view(labels.size(0), n, -1)
+        grams = labels[:, 1:].repeat(1, n).view(labels.size(0), n, -1)
+
+    for i in range(1, grams.size(1)):
+        grams[:, i] |= torch.roll(grams[:, i - 1], 1)
 
-    for i in range(1, grams.size(1)):        
-        grams[:,i] |= torch.roll(grams[:,i-1], 1)
-        
         if ignore_idx is not None:
-            grams[:,i][grams[:,i] < 0] = ignore_idx
+            grams[:, i][grams[:, i] < 0] = ignore_idx
+
+    for i in range(1, grams.size(1)):
+        grams[:, i, :i] = ignore_idx
 
-    for i in range(1, grams.size(1)):                     
-        grams[:,i,:i] = ignore_idx 
-        
     if ignore_idx is not None:
         v = (torch.zeros(grams.size(0) * n, dtype=torch.long) + ignore_idx).view(grams.size(0), n, 1).to(labels.device)
         grams = torch.cat([v, grams], dim=2)
 
-    return grams.permute(0,2,1).contiguous()
+    return grams.permute(0, 2, 1).contiguous()
diff --git a/transformers_lightning/utils/readers.py b/transformers_lightning/utils/readers.py
index 48c65f1..fa188bb 100644
--- a/transformers_lightning/utils/readers.py
+++ b/transformers_lightning/utils/readers.py
@@ -18,6 +18,7 @@ def load_yaml(filename: str, to_namespace: bool = True) -> Union[Dict, Namespace
         return Namespace(**res)
     return res
 
+
 def load_json(filename: str, to_namespace: bool = True) -> Union[Dict, Namespace]:
     r"""
     Load config from json file from disk
@@ -29,21 +30,21 @@ def load_json(filename: str, to_namespace: bool = True) -> Union[Dict, Namespace
         return Namespace(**res)
     return res
 
+
 def dump_json(filename: str, data: Dict, complain: bool = False) -> None:
     r"""
     Save json to file.
     """
     if os.path.isfile(filename):
         if complain:
-            raise ValueError(
-                f"File {filename} does already exist!"
-            )
+            raise ValueError(f"File {filename} does already exist!")
         else:
             logger.warn(f"Overwriting {filename} file!")
 
     with open(filename, 'w') as outfile:
         json.dump(data, outfile)
 
+
 def strip_lines(iterable: Iterable[str]) -> Generator:
     r"""
     Remove blank lines from iterable over strings and return new generator.


Test patch:
diff --git a/tests/adapters/test_csv_adapter.py b/tests/adapters/test_csv_adapter.py
index b6e6c99..2145097 100644
--- a/tests/adapters/test_csv_adapter.py
+++ b/tests/adapters/test_csv_adapter.py
@@ -8,12 +8,23 @@ class ExampleAdapter(CSVAdapter):
     def preprocess_line(self, line: list) -> list:
         return [int(line[0]), int(line[1]), int(line[2]), line[3], line[4], eval(line[5])]
 
+
 # Test iter dataset work correctly with dp
 @pytest.mark.parametrize(
     ["test_file", "line_number", "expected_line", "total_len"], [
-        ["tests/test_data/test1.tsv", 3, [3, 33115, 4156, "How can I reduce tummy fat?", "How do I reduce tummy?", True], 40],
-        ["tests/test_data/test2.tsv", 54, [54, 522296, 3013, "What's it like to take a daily caffeine pill?", "How can I improve digestion?", False], 96],
-        ["tests/test_data/test3.tsv", 13, [13, 351690, 490931, "How do I withdraw at Quora?", "How can I remove my account from Quora?", True], 17]
+        [
+            "tests/test_data/test1.tsv", 3,
+            [3, 33115, 4156, "How can I reduce tummy fat?", "How do I reduce tummy?", True], 40
+        ],
+        [
+            "tests/test_data/test2.tsv", 54,
+            [54, 522296, 3013, "What's it like to take a daily caffeine pill?", "How can I improve digestion?", False],
+            96
+        ],
+        [
+            "tests/test_data/test3.tsv", 13,
+            [13, 351690, 490931, "How do I withdraw at Quora?", "How can I remove my account from Quora?", True], 17
+        ]
     ]
 )
 def test_datamodule_gpu_dp(test_file, line_number, expected_line, total_len):
@@ -22,9 +33,5 @@ def test_datamodule_gpu_dp(test_file, line_number, expected_line, total_len):
     adapter = ExampleAdapter(hparams, test_file)
     data = [adapter.preprocess_line(line) for line in iter(adapter)]
 
-    assert data[line_number] == expected_line, (
-        f"Expected: {expected_line}, got: {data[line_number]}"
-    )
+    assert data[line_number] == expected_line, (f"Expected: {expected_line}, got: {data[line_number]}")
     assert len(data) == total_len, f"Expected: {total_len}, got: {len(data)}"
-
-
diff --git a/tests/adapters/test_transformers_adapter.py b/tests/adapters/test_transformers_adapter.py
index 44f32e9..63fdb26 100644
--- a/tests/adapters/test_transformers_adapter.py
+++ b/tests/adapters/test_transformers_adapter.py
@@ -9,31 +9,47 @@ class ExampleAdapter(TransformersAdapter):
 
     def preprocess_line(self, line: list) -> list:
         line = [int(line[0]), int(line[1]), int(line[2]), line[3], line[4], eval(line[5])]
-    
-        results = self.tokenizer.encode_plus(
-            line[3], line[4],
-            truncation=True
-        )
+
+        results = self.tokenizer.encode_plus(line[3], line[4], truncation=True)
         results['words_tails'] = self._convert_ids_to_word_tails(results['input_ids'])
         return results
 
+
 # Test iter dataset work correctly with dp
 @pytest.mark.parametrize(
     ["test_file", "line_number", "expected_ids", "expected_word_tails"], [
         [
             "tests/test_data/test1.tsv", 3,
-            [101, 1731, 1169, 146, 4851, 189, 1818, 4527, 7930, 136, 102, 1731, 1202, 146, 4851, 189, 1818, 4527, 136, 102],
-            [False, False, False, False, False, False, True, True, False, False, False, False, False, False, False, False, True, True, False, False]
+            [
+                101, 1731, 1169, 146, 4851, 189, 1818, 4527, 7930, 136, 102, 1731, 1202, 146, 4851, 189, 1818, 4527,
+                136, 102
+            ],
+            [
+                False, False, False, False, False, False, True, True, False, False, False, False, False, False, False,
+                False, True, True, False, False
+            ]
         ],
         [
             "tests/test_data/test2.tsv", 54,
-            [101, 1327, 112, 188, 1122, 1176, 1106, 1321, 170, 3828, 11019, 15475, 2042, 21822, 136, 102, 1731, 1169, 146, 4607, 11902, 2556, 1988, 136, 102],
-            [False, False, False, False, False, False, False, False, False, False, False, True, True, False, False, False, False, False, False, False, False, True, True, False, False]
+            [
+                101, 1327, 112, 188, 1122, 1176, 1106, 1321, 170, 3828, 11019, 15475, 2042, 21822, 136, 102, 1731, 1169,
+                146, 4607, 11902, 2556, 1988, 136, 102
+            ],
+            [
+                False, False, False, False, False, False, False, False, False, False, False, True, True, False, False,
+                False, False, False, False, False, False, True, True, False, False
+            ]
         ],
         [
             "tests/test_data/test3.tsv", 13,
-            [101, 1731, 1202, 146, 10570, 1120, 154, 11848, 1611, 136, 102, 1731, 1169, 146, 5782, 1139, 3300, 1121, 154, 11848, 1611, 136, 102],
-            [False, False, False, False, False, False, False, True, True, False, False, False, False, False, False, False, False, False, False, True, True, False, False]
+            [
+                101, 1731, 1202, 146, 10570, 1120, 154, 11848, 1611, 136, 102, 1731, 1169, 146, 5782, 1139, 3300, 1121,
+                154, 11848, 1611, 136, 102
+            ],
+            [
+                False, False, False, False, False, False, False, True, True, False, False, False, False, False, False,
+                False, False, False, False, True, True, False, False
+            ]
         ]
     ]
 )
diff --git a/tests/callbacks/test_ckpt_callback.py b/tests/callbacks/test_ckpt_callback.py
index 6dee0ed..655279c 100644
--- a/tests/callbacks/test_ckpt_callback.py
+++ b/tests/callbacks/test_ckpt_callback.py
@@ -21,33 +21,38 @@ class ExampleDataModule(SuperDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
-        self.valid_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
-        self.test_adapter = [ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer) for _ in range(2)]
+        self.train_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
+        self.valid_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
+        self.test_adapter = [
+            ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
+            for _ in range(2)
+        ]
 
 
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
     ["epochs", "accumulate_grad_batches", "batch_size", "callback_interval", "val_callback", "expected_results"], [
-    [2,         3,                         4,            3,                   False,         ["hparams.json",
-                                                                                              "ckpt_epoch_0_step_3",
-                                                                                              "ckpt_epoch_0_step_6",
-                                                                                              "ckpt_epoch_0_step_8",
-                                                                                              "ckpt_epoch_1_step_9",
-                                                                                              "ckpt_epoch_1_step_12",
-                                                                                              "ckpt_epoch_1_step_15",
-                                                                                              "ckpt_epoch_1_step_16_final"]],
-    [1,         2,                         5,            6,                   False,         ["hparams.json",
-                                                                                              "ckpt_epoch_0_step_6",
-                                                                                              "ckpt_epoch_0_step_10_final"]],
-    [1,         2,                         5,            6,                   True,          ["hparams.json",
-                                                                                              "ckpt_epoch_0_step_1",
-                                                                                              "ckpt_epoch_0_step_3",
-                                                                                              "ckpt_epoch_0_step_5",
-                                                                                              "ckpt_epoch_0_step_6",
-                                                                                              "ckpt_epoch_0_step_8",
-                                                                                              "ckpt_epoch_0_step_10_final"]],
-])
+        [
+            2, 3, 4, 3, False,
+            [
+                "hparams.json", "ckpt_epoch_0_step_3", "ckpt_epoch_0_step_6", "ckpt_epoch_0_step_8",
+                "ckpt_epoch_1_step_9", "ckpt_epoch_1_step_12", "ckpt_epoch_1_step_15", "ckpt_epoch_1_step_16_final"
+            ]
+        ],
+        [1, 2, 5, 6, False, ["hparams.json", "ckpt_epoch_0_step_6", "ckpt_epoch_0_step_10_final"]],
+        [
+            1, 2, 5, 6, True,
+            [
+                "hparams.json", "ckpt_epoch_0_step_1", "ckpt_epoch_0_step_3", "ckpt_epoch_0_step_5",
+                "ckpt_epoch_0_step_6", "ckpt_epoch_0_step_8", "ckpt_epoch_0_step_10_final"
+            ]
+        ],
+    ]
+)
 def test_datamodule_cpu(epochs, accumulate_grad_batches, batch_size, callback_interval, val_callback, expected_results):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -86,7 +91,7 @@ def test_datamodule_cpu(epochs, accumulate_grad_batches, batch_size, callback_in
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, test_number=2, tokenizer=tokenizer)
@@ -98,6 +103,3 @@ def test_datamodule_cpu(epochs, accumulate_grad_batches, batch_size, callback_in
     listing = os.listdir(folder)
     shutil.rmtree(hparams.output_dir)
     assert set(listing) == set(expected_results), f"{listing} vs {set(expected_results)}"
-
-    
-
diff --git a/tests/datamodules/test_datamodule_cpu.py b/tests/datamodules/test_datamodule_cpu.py
index 3dadf61..0d44408 100644
--- a/tests/datamodules/test_datamodule_cpu.py
+++ b/tests/datamodules/test_datamodule_cpu.py
@@ -12,41 +12,43 @@
 
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "gpus", "epochs"], [
-    
+    ["ds_type", "num_workers", "gpus", "epochs"],
+    [
+
     # ITER dataset
     # test different num_workers in single node on cpu
-    ['iter',     0,             0,   1],
-    ['iter',     1,             0,   1],
-    ['iter',     2,             0,   1],
-    ['iter',     n_cpus,        0,   1],
-    
+        ['iter', 0, 0, 1],
+        ['iter', 1, 0, 1],
+        ['iter', 2, 0, 1],
+        ['iter', n_cpus, 0, 1],
+
     # num_workers through epochs
-    ['iter',     0,             0,   1],
-    ['iter',     0,             0,   4],
-    ['iter',     0,             0,   2],
-    ['iter',     0,             0,   10],
-    ['iter',     2,             0,   1],
-    ['iter',     2,             0,   2],
-    ['iter',     2,             0,   4],
+        ['iter', 0, 0, 1],
+        ['iter', 0, 0, 4],
+        ['iter', 0, 0, 2],
+        ['iter', 0, 0, 10],
+        ['iter', 2, 0, 1],
+        ['iter', 2, 0, 2],
+        ['iter', 2, 0, 4],
 
     # MAP dataset
     # test different num_workers in single node on cpu
-    ['map',     0,             0,   1],
-    ['map',     1,             0,   1],
-    ['map',     2,             0,   1],
-    ['map',     n_cpus,        0,   1],
-    
+        ['map', 0, 0, 1],
+        ['map', 1, 0, 1],
+        ['map', 2, 0, 1],
+        ['map', n_cpus, 0, 1],
+
     # num_workers through epochs
-    ['map',     0,             0,   1],
-    ['map',     0,             0,   2],
-    ['map',     0,             0,   4],
-    ['map',     2,             0,   1],
-    ['map',     2,             0,   2],
-    ['map',     2,             0,   4],
-])
+        ['map', 0, 0, 1],
+        ['map', 0, 0, 2],
+        ['map', 0, 0, 4],
+        ['map', 2, 0, 1],
+        ['map', 2, 0, 2],
+        ['map', 2, 0, 4],
+    ]
+)
 def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
-    
+
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
     hparams = Namespace(
@@ -75,7 +77,7 @@ def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, test_number=1, tokenizer=tokenizer)
diff --git a/tests/datamodules/test_datamodule_ddp.py b/tests/datamodules/test_datamodule_ddp.py
index 5945592..9f00d79 100644
--- a/tests/datamodules/test_datamodule_ddp.py
+++ b/tests/datamodules/test_datamodule_ddp.py
@@ -14,28 +14,30 @@
 # Test iter dataset work correctly with dp
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs", "dataset_idx"], [
+    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs", "dataset_idx"],
+    [
 
     # ITER dataset
     # num_workers with ddp
-    ['iter',     0,             'ddp',                  2,      1,      1],
-    ['iter',     1,             'ddp',                  2,      2,      1],
-    ['iter',     2,             'ddp',                  2,      2,      1],
-    ['iter',     0,             'ddp',                  2,      1,      2],
-    ['iter',     n_cpus,        'ddp',                  2,      5,     2],
+        ['iter', 0, 'ddp', 2, 1, 1],
+        ['iter', 1, 'ddp', 2, 2, 1],
+        ['iter', 2, 'ddp', 2, 2, 1],
+        ['iter', 0, 'ddp', 2, 1, 2],
+        ['iter', n_cpus, 'ddp', 2, 5, 2],
 
     # MAP dataset
     # num_workers with ddp
-    ['map',     0,             'ddp',                  2,      2,       1],
-    ['map',     1,             'ddp',                  2,      2,       1],
-    ['map',     2,             'ddp',                  2,      2,       1],
-    ['map',     0,             'ddp',                  2,      1,       2],
-    ['map',     n_cpus,        'ddp',                  2,      5,      3],
-])
+        ['map', 0, 'ddp', 2, 2, 1],
+        ['map', 1, 'ddp', 2, 2, 1],
+        ['map', 2, 'ddp', 2, 2, 1],
+        ['map', 0, 'ddp', 2, 1, 2],
+        ['map', n_cpus, 'ddp', 2, 5, 3],
+    ]
+)
 def test_datamodule_gpu_ddp_only(ds_type, num_workers, distributed_backend, gpus, epochs, dataset_idx):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
-    time.sleep(2) # sleep for 5 second to be sure area is clean
+    time.sleep(2)    # sleep for 5 second to be sure area is clean
 
     hparams = Namespace(
         batch_size=4,
@@ -64,7 +66,7 @@ def test_datamodule_gpu_ddp_only(ds_type, num_workers, distributed_backend, gpus
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams, do_ids_check=(ds_type != 'map'))    
+    model = SimpleTransformerLikeModel(hparams, do_ids_check=(ds_type != 'map'))
 
     # Datasets
     datamodule = ExampleDataModule(hparams, test_number=1, tokenizer=tokenizer)
diff --git a/tests/datamodules/test_datamodule_dp.py b/tests/datamodules/test_datamodule_dp.py
index 3efe8b1..b2d4768 100644
--- a/tests/datamodules/test_datamodule_dp.py
+++ b/tests/datamodules/test_datamodule_dp.py
@@ -15,36 +15,36 @@
 # Test iter dataset work correctly with dp
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs"], [
-    
+    ["ds_type", "num_workers", "distributed_backend", "gpus", "epochs"],
+    [
+
     # ITER dataset
     # num_workers with dp
-    ['iter',     0,             'dp',                   2,      2],
-    ['iter',     1,             'dp',                   2,      2],
-    ['iter',     2,             'dp',                   2,      2],
-    ['iter',     n_cpus,        'dp',                   2,      2],
-
-    ['iter',     0,             'dp',                   2,      1],
-    ['iter',     1,             'dp',                   2,      2],
-    ['iter',     2,             'dp',                   2,      4],
-    ['iter',     n_cpus,        'dp',                   2,      5],
+        ['iter', 0, 'dp', 2, 2],
+        ['iter', 1, 'dp', 2, 2],
+        ['iter', 2, 'dp', 2, 2],
+        ['iter', n_cpus, 'dp', 2, 2],
+        ['iter', 0, 'dp', 2, 1],
+        ['iter', 1, 'dp', 2, 2],
+        ['iter', 2, 'dp', 2, 4],
+        ['iter', n_cpus, 'dp', 2, 5],
 
     # MAP dataset
     # num_workers with dp
-    ['map',     0,             'dp',                   2,      2],
-    ['map',     1,             'dp',                   2,      2],
-    ['map',     2,             'dp',                   2,      2],
-    ['map',     n_cpus,        'dp',                   2,      2],
-
-    ['map',     0,             'dp',                   2,      1],
-    ['map',     1,             'dp',                   2,      2],
-    ['map',     2,             'dp',                   2,      4],
-    ['map',     n_cpus,        'dp',                   2,      5]
-])
+        ['map', 0, 'dp', 2, 2],
+        ['map', 1, 'dp', 2, 2],
+        ['map', 2, 'dp', 2, 2],
+        ['map', n_cpus, 'dp', 2, 2],
+        ['map', 0, 'dp', 2, 1],
+        ['map', 1, 'dp', 2, 2],
+        ['map', 2, 'dp', 2, 4],
+        ['map', n_cpus, 'dp', 2, 5]
+    ]
+)
 def test_datamodule_gpu_dp(ds_type, num_workers, distributed_backend, gpus, epochs):
 
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"    
-    time.sleep(2) # sleep for 5 second to be sure area is clean
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+    time.sleep(2)    # sleep for 5 second to be sure area is clean
 
     hparams = Namespace(
         batch_size=4,
@@ -73,11 +73,10 @@ def test_datamodule_gpu_dp(ds_type, num_workers, distributed_backend, gpus, epoc
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, test_number=1, tokenizer=tokenizer)
 
     model.datamodule = datamodule
     trainer.fit(model, datamodule=datamodule)
-
diff --git a/tests/datamodules/test_datamodule_val_test.py b/tests/datamodules/test_datamodule_val_test.py
index 4ad6014..52ddfeb 100644
--- a/tests/datamodules/test_datamodule_val_test.py
+++ b/tests/datamodules/test_datamodule_val_test.py
@@ -15,35 +15,44 @@ class ExampleDataModule(datamodules.SuperDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
-        self.valid_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
-        self.test_adapter = [ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer) for _ in range(2)]
+        self.train_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
+        self.valid_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
+        self.test_adapter = [
+            ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
+            for _ in range(2)
+        ]
 
 
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["ds_type", "num_workers", "gpus", "epochs"], [
-    
+    ["ds_type", "num_workers", "gpus", "epochs"],
+    [
+
     # ITER dataset
     # test different num_workers in single node on cpu
-    ['iter',     0,             0,   1],
-    ['iter',     n_cpus,        0,   1],
-    
+        ['iter', 0, 0, 1],
+        ['iter', n_cpus, 0, 1],
+
     # num_workers through epochs
-    ['iter',     0,             0,   1],
-    ['iter',     0,             0,   10],
+        ['iter', 0, 0, 1],
+        ['iter', 0, 0, 10],
 
     # MAP dataset
     # test different num_workers in single node on cpu
-    ['map',     0,             0,   1],
-    ['map',     n_cpus,        0,   1],
-    
+        ['map', 0, 0, 1],
+        ['map', n_cpus, 0, 1],
+
     # num_workers through epochs
-    ['map',     0,             0,   1],
-    ['map',     2,             0,   1],
-])
+        ['map', 0, 0, 1],
+        ['map', 2, 0, 1],
+    ]
+)
 def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
-    
+
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
     hparams = Namespace(
@@ -72,7 +81,7 @@ def test_datamodule_cpu(ds_type, num_workers, gpus, epochs):
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, test_number=1, tokenizer=tokenizer)
diff --git a/tests/datamodules/test_utils.py b/tests/datamodules/test_utils.py
index 2d157c6..d0373f8 100644
--- a/tests/datamodules/test_utils.py
+++ b/tests/datamodules/test_utils.py
@@ -1,6 +1,5 @@
 import torch
-from transformers.models.bert.modeling_bert import (BertConfig,
-                                        BertForSequenceClassification)
+from transformers.models.bert.modeling_bert import (BertConfig, BertForSequenceClassification)
 from transformers import AdamW
 from transformers_lightning.adapters.csv_adapter import CSVAdapter
 from transformers_lightning import datamodules, models
@@ -15,19 +14,21 @@ def __init__(self, hparams, do_ids_check=True):
         # super light BERT model
         config = BertConfig(hidden_size=12, num_hidden_layers=1, num_attention_heads=1, intermediate_size=12)
         self.model = BertForSequenceClassification(config)
-    
+
     def configure_optimizers(self):
         return AdamW(self.model.parameters())
 
     def training_step(self, batch, batch_idx):
 
         if self.trainer.distributed_backend == "ddp":
-            print(f"ID {torch.distributed.get_rank()}/{torch.distributed.get_world_size()} processing ids: {batch['ids']}")
+            print(
+                f"ID {torch.distributed.get_rank()}/{torch.distributed.get_world_size()} processing ids: {batch['ids']}"
+            )
 
         batch['labels'] = batch['labels'].to(dtype=torch.long)
         kwargs = {k: batch[k] for k in ["input_ids", "attention_mask", "token_type_ids", "labels"]}
         results = self(**kwargs)
-        return { 'loss': results.loss, 'ids': batch['ids'] }
+        return {'loss': results.loss, 'ids': batch['ids']}
 
     def training_step_end(self, batch_parts):
         batch_parts['loss'] = torch.sum(batch_parts['loss'])
@@ -48,26 +49,23 @@ def training_epoch_end(self, outputs):
             print(f"ID {torch.distributed.get_rank()}/{torch.distributed.get_world_size()} ALL ids: {ids}")
 
         try:
-            received = torch.zeros((len(self.datamodule.train_dataset),)).to(dtype=bool)
+            received = torch.zeros((len(self.datamodule.train_dataset), )).to(dtype=bool)
         except TypeError:
             if self.trainer.distributed_backend == "ddp":
-                expected_len = (
-                    self.datamodule.train_dataset.length // torch.distributed.get_world_size()
-                ) * torch.distributed.get_world_size()
+                expected_len = (self.datamodule.train_dataset.length //
+                                torch.distributed.get_world_size()) * torch.distributed.get_world_size()
             else:
                 expected_len = self.datamodule.train_dataset.length
-            received = torch.zeros((expected_len,)).to(dtype=bool)
+            received = torch.zeros((expected_len, )).to(dtype=bool)
         received[ids] = True
 
         if self.do_ids_check:
             # assert no duplicate element received
-            assert len(set(ids.tolist())) == len(ids.tolist()), (
-                f"Received {len(ids.tolist())} ids but only {len(set(ids.tolist()))} are unique: {ids}"
-            )
+            assert len(set(ids.tolist())) == len(
+                ids.tolist()
+            ), (f"Received {len(ids.tolist())} ids but only {len(set(ids.tolist()))} are unique: {ids}")
             # assert all elements received
-            assert all(received), (
-                f"({self.trainer.max_steps}) Received not all {len(received)} ids: {received}"
-            )
+            assert all(received), (f"({self.trainer.max_steps}) Received not all {len(received)} ids: {received}")
 
     def validation_step(self, batch, batch_idx):
         batch['labels'] = batch['labels'].to(dtype=torch.long)
@@ -91,14 +89,15 @@ def __init__(self, *args, tokenizer=None, **kwargs):
     def preprocess_line(self, line: list) -> list:
 
         results = self.tokenizer.encode_plus(
-            line[3], line[4],
+            line[3],
+            line[4],
             add_special_tokens=True,
             padding='max_length',
             max_length=self.hparams.max_sequence_length,
             truncation=True
         )
 
-        res = { **results, 'ids': int(line[0]), 'labels': line[5].lower().strip() == "true" }
+        res = {**results, 'ids': int(line[0]), 'labels': line[5].lower().strip() == "true"}
         return res
 
 
@@ -106,4 +105,6 @@ class ExampleDataModule(datamodules.SuperDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)
+        self.train_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
diff --git a/tests/datasets/test_start_from_steps.py b/tests/datasets/test_start_from_steps.py
index 020025e..521da6f 100644
--- a/tests/datasets/test_start_from_steps.py
+++ b/tests/datasets/test_start_from_steps.py
@@ -13,6 +13,7 @@
 n_cpus = multiprocessing.cpu_count()
 N = 100
 
+
 class SimpleTransformerLikeModel(models.TransformersModel):
 
     def __init__(self, hparams):
@@ -25,7 +26,7 @@ def configure_optimizers(self):
 
     def training_step(self, batch, batch_idx):
         results = self.lin(batch['data']).mean()
-        return { 'loss': results, 'ids': batch['ids'] }
+        return {'loss': results, 'ids': batch['ids']}
 
     def training_epoch_end(self, outputs):
         self.all_received_ids = torch.cat([output['ids'] for output in outputs])
@@ -48,7 +49,7 @@ def __iter__(self):
             yield (i, [1.0] * 10, torch.LongTensor(1))
 
     def preprocess_line(self, line: list) -> list:
-        res = { 'ids': line[0], 'data': line[1], 'labels': line[2] }
+        res = {'ids': line[0], 'data': line[1], 'labels': line[2]}
         return res
 
 
@@ -56,16 +57,17 @@ class ExampleDataModule(datamodules.SuperDataModule):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams)        
+        self.train_adapter = ExampleAdapter(self.hparams)
 
 
 # Test if skip_steps works correctly
 @pytest.mark.parametrize(
     ["max_epochs", "accumulate_grad_batches", "batch_size", "skip", "expected_ids"], [
-    [1,             1,                         4,            4,         list(range(16, 100))],
-    [1,             3,                         8,            3,         list(range(72, 100))],
-    [1,             2,                         5,           5,          list(range(50, 100))],
-])
+        [1, 1, 4, 4, list(range(16, 100))],
+        [1, 3, 8, 3, list(range(72, 100))],
+        [1, 2, 5, 5, list(range(50, 100))],
+    ]
+)
 def test_skip_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, skip, expected_ids):
 
     hparams = Namespace(
@@ -90,25 +92,26 @@ def test_skip_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, skip, e
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams)
 
     # Train!
     trainer.fit(model, datamodule=datamodule)
-    assert sorted(model.all_received_ids) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
-
+    assert sorted(model.all_received_ids
+                 ) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
 
 
 # Test if skip_steps works correctly in ddp
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
     ["max_epochs", "accumulate_grad_batches", "batch_size", "skip", "expected_ids"], [
-    [1,             1,                         4,            4,         list(range(32, 100))],
-    [1,             3,                         8,            2,         list(range(96, 100))],
-    [1,             2,                         7,            2,         list(range(28, 100))],
-])
+        [1, 1, 4, 4, list(range(32, 100))],
+        [1, 3, 8, 2, list(range(96, 100))],
+        [1, 2, 7, 2, list(range(28, 100))],
+    ]
+)
 def test_skip_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, skip, expected_ids):
 
     hparams = Namespace(
@@ -134,11 +137,12 @@ def test_skip_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, skip, e
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams)
 
     # Train!
     trainer.fit(model, datamodule=datamodule)
-    assert sorted(model.all_received_ids) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
\ No newline at end of file
+    assert sorted(model.all_received_ids
+                 ) == sorted(expected_ids), f"{sorted(model.all_received_ids)} != {sorted(expected_ids)}"
diff --git a/tests/language_models/test_mlm.py b/tests/language_models/test_mlm.py
index 6cab5af..ca18af9 100644
--- a/tests/language_models/test_mlm.py
+++ b/tests/language_models/test_mlm.py
@@ -8,14 +8,29 @@
 tok = BertTokenizer.from_pretrained('bert-base-cased')
 mlm = MaskedLanguageModeling(tok, whole_word_masking=True)
 
+
 # Test iter dataset work correctly
-@pytest.mark.parametrize(["seed", "sentence", "masking"],
-    [
+@pytest.mark.parametrize(
+    ["seed", "sentence", "masking"], [
         [0, "how are you man?", [IGNORE_IDX, IGNORE_IDX, 1132, 1128, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX]],
-        [1, "The quick brown fox jumps over the lazy dog", [IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, 17594, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX]],
-        [8, "Be or not to be a superstar",  [IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, 7688, 10058, IGNORE_IDX]],
-        [3, "", [IGNORE_IDX, IGNORE_IDX]],
-        [4, "share silence or say what you think?", [IGNORE_IDX, IGNORE_IDX, 3747, IGNORE_IDX, 1474, 1184, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX]]
+        [
+            1, "The quick brown fox jumps over the lazy dog",
+            [
+                IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, 17594, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX,
+                IGNORE_IDX, IGNORE_IDX
+            ]
+        ],
+        [
+            8, "Be or not to be a superstar",
+            [
+                IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, 7688, 10058,
+                IGNORE_IDX
+            ]
+        ], [3, "", [IGNORE_IDX, IGNORE_IDX]],
+        [
+            4, "share silence or say what you think?",
+            [IGNORE_IDX, IGNORE_IDX, 3747, IGNORE_IDX, 1474, 1184, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX, IGNORE_IDX]
+        ]
     ]
 )
 def test_datamodule_cpu(seed, sentence, masking):
diff --git a/tests/language_models/test_rts.py b/tests/language_models/test_rts.py
index 3b0e183..15680ae 100644
--- a/tests/language_models/test_rts.py
+++ b/tests/language_models/test_rts.py
@@ -6,18 +6,26 @@
 from transformers_lightning.language_modeling.random_token_substitution import RandomTokenSubstitution
 from transformers_lightning.language_modeling.utils import whole_word_tails_mask
 
-
 tok = BertTokenizer.from_pretrained('bert-base-cased')
 rts = RandomTokenSubstitution(tok, whole_word_swapping=True)
 
+
 # Test iter dataset work correctly
-@pytest.mark.parametrize(["seed", "sentence", "masking", "new_ids"],
-    [
+@pytest.mark.parametrize(
+    ["seed", "sentence", "masking", "new_ids"], [
         [0, "how are you man?", [IGNORE_IDX, 0, 1, 1, 0, 0, IGNORE_IDX], [101, 1293, 28135, 17785, 1299, 136, 102]],
-        [1, "The quick brown fox jumps over the lazy dog", [IGNORE_IDX, 0, 0, 0, 1, 0, 0, 0, 0, 0, IGNORE_IDX], [101, 1109, 3613, 3058, 18156, 15457, 1166, 1103, 16688, 3676, 102]],
-        [8, "Be or not to be a superstar",  [IGNORE_IDX, 0, 0, 0, 0, 0, 0, 1, 1, IGNORE_IDX], [101, 4108, 1137, 1136, 1106, 1129, 170, 4667, 20540, 102]],
-        [3, "", [IGNORE_IDX, IGNORE_IDX], [101, 102]],
-        [4, "share silence or say what you think?", [IGNORE_IDX, 0, 1, 0, 1, 1, 0, 0, 0, IGNORE_IDX], [101, 2934, 13287, 1137, 24438, 26891, 1128, 1341, 136, 102]]
+        [
+            1, "The quick brown fox jumps over the lazy dog", [IGNORE_IDX, 0, 0, 0, 1, 0, 0, 0, 0, 0, IGNORE_IDX],
+            [101, 1109, 3613, 3058, 18156, 15457, 1166, 1103, 16688, 3676, 102]
+        ],
+        [
+            8, "Be or not to be a superstar", [IGNORE_IDX, 0, 0, 0, 0, 0, 0, 1, 1, IGNORE_IDX],
+            [101, 4108, 1137, 1136, 1106, 1129, 170, 4667, 20540, 102]
+        ], [3, "", [IGNORE_IDX, IGNORE_IDX], [101, 102]],
+        [
+            4, "share silence or say what you think?", [IGNORE_IDX, 0, 1, 0, 1, 1, 0, 0, 0, IGNORE_IDX],
+            [101, 2934, 13287, 1137, 24438, 26891, 1128, 1341, 136, 102]
+        ]
     ]
 )
 def test_datamodule_cpu(seed, sentence, masking, new_ids):
diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py
index 11f688c..a73adcf 100644
--- a/tests/metrics/test_metrics.py
+++ b/tests/metrics/test_metrics.py
@@ -2,7 +2,7 @@
 import torch
 
 from transformers_lightning.metrics.retrieval import (
-    MeanReciprocalRank, MeanAveragePrecision, PrecisionAtK, RecallAtK, HitRateAtK
+    MeanReciprocalRank, MeanAveragePrecision, Precision, Recall, HitRateAtK
 )
 
 
@@ -12,44 +12,32 @@
         [
             torch.tensor([0, 0, 0, 1, 1, 1, 1]),
             torch.tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2]),
-            torch.tensor([False, False, True, False, True, False, False]),
-            "skip",
-            1,
-            0.75, 0.75, 0.5, 0.5, 0.5
+            torch.tensor([False, False, True, False, True, False, False]), "skip", 1, 0.75, 0.75, 0.5, 0.5, 0.5
         ],
         [
-            torch.tensor([  0,  0,  0,  0,  1,  1,  1,  1,  1]),
-            torch.tensor([  0.2,0.1,0.3,0.1,0.7,0.4,0.3,0.2,0.1]),
-            torch.tensor([  1,  0,  1,  0,  0,  1,  0,  0,  0]),
-            "skip",
-            1,
-            0.75, 0.75, 0.5, 0.25, 0.5
+            torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 1]),
+            torch.tensor([0.2, 0.1, 0.3, 0.1, 0.7, 0.4, 0.3, 0.2, 0.1]),
+            torch.tensor([1, 0, 1, 0, 0, 1, 0, 0, 0]), "skip", 1, 0.75, 0.75, 0.5, 0.25, 0.5
         ],
         [
-            torch.tensor([  0,  0,  0,   0,  1,  1,  1,  1,  1]),
-            torch.tensor([  0.2,0.1,-0.3,0.1,0.7,0.4,0.3,0.2,0.1]),
-            torch.tensor([  1,  0,  1,   0,  0,  0,  0,  0,  0]),
-            "positive",
-            2,
-            1.0, 0.875, 0.75, 0.75, 1.0
+            torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 1]),
+            torch.tensor([0.2, 0.1, -0.3, 0.1, 0.7, 0.4, 0.3, 0.2, 0.1]),
+            torch.tensor([1, 0, 1, 0, 0, 0, 0, 0, 0]), "pos", 2, 1.0, 0.875, 0.75, 0.75, 1.0
         ],
         [
-            torch.tensor([  0,  0,  0,   0,  1,  1,  1,  1,  1]),
-            torch.tensor([  0.2,0.1,-0.3,0.1,0.7,0.4,0.3,0.2,0.1]),
-            torch.tensor([  1,  0,  1,   0,  0,  0,  0,  0,  0]),
-            "negative",
-            1,
-            0.5, 0.375, 0.5, 0.25, 0.5
+            torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 1]),
+            torch.tensor([0.2, 0.1, -0.3, 0.1, 0.7, 0.4, 0.3, 0.2, 0.1]),
+            torch.tensor([1, 0, 1, 0, 0, 0, 0, 0, 0]), "neg", 1, 0.5, 0.375, 0.5, 0.25, 0.5
         ],
     ]
 )
 def test_metrics(ids, predictions, labels, empty_docs, k, mrr, map, precision, recall, hit_rate):
 
-    _mrr = MeanReciprocalRank(empty_documents=empty_docs)
-    _map = MeanAveragePrecision(empty_documents=empty_docs)
-    _prec = PrecisionAtK(k=k, empty_documents=empty_docs)
-    _rec = RecallAtK(k=k, empty_documents=empty_docs)
-    _hr = HitRateAtK(k=k, empty_documents=empty_docs)
+    _mrr = MeanReciprocalRank(query_without_relevant_docs=empty_docs)
+    _map = MeanAveragePrecision(query_without_relevant_docs=empty_docs)
+    _prec = Precision(k=k, query_without_relevant_docs=empty_docs)
+    _rec = Recall(k=k, query_without_relevant_docs=empty_docs)
+    _hr = HitRateAtK(k=k, query_without_relevant_docs=empty_docs)
 
     _mrr = _mrr(ids, predictions, labels)
     _map = _map(ids, predictions, labels)
diff --git a/tests/models/test_fix_max_steps.py b/tests/models/test_fix_max_steps.py
index b703f7e..2818f17 100644
--- a/tests/models/test_fix_max_steps.py
+++ b/tests/models/test_fix_max_steps.py
@@ -10,11 +10,11 @@
 import torch
 from transformers_lightning import datamodules, models, adapters, utils
 from transformers import BertTokenizer
-from transformers.models.bert.modeling_bert import (BertConfig,
-                                        BertForSequenceClassification)
+from transformers.models.bert.modeling_bert import (BertConfig, BertForSequenceClassification)
 
 n_cpus = multiprocessing.cpu_count()
 
+
 class SimpleTransformerLikeModel(models.TransformersModel):
 
     def __init__(self, hparams):
@@ -43,13 +43,14 @@ def __init__(self, *args, tokenizer=None, **kwargs):
     def preprocess_line(self, line: list) -> list:
 
         results = self.tokenizer.encode_plus(
-            line[3], line[4],
+            line[3],
+            line[4],
             add_special_tokens=True,
             padding='max_length',
             max_length=self.hparams.max_sequence_length,
             truncation=True
         )
-        res = { **results, 'labels': line[5].strip().lower() == "true" }
+        res = {**results, 'labels': line[5].strip().lower() == "true"}
         return res
 
 
@@ -57,23 +58,24 @@ class ExampleDataModule(datamodules.SuperDataModule):
 
     def __init__(self, *args, test_number=1, tokenizer=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.train_adapter = ExampleAdapter(self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer)        
+        self.train_adapter = ExampleAdapter(
+            self.hparams, f"tests/test_data/test{test_number}.tsv", delimiter="\t", tokenizer=tokenizer
+        )
 
 
 # Test if max_steps fix works correctly
 @pytest.mark.parametrize(
     ["max_epochs", "accumulate_grad_batches", "batch_size", "expected_max_steps", "ds_type"], [
-
-    [1,             1,                         4,            10,         'map'],
-    [1,             3,                         8,            2,          'map'],
-    [4,             2,                         12,           8,          'map'],
-    [4,             4,                         16,           4,          'map'],
-
-    [1,             1,                         4,            10,         'iter'],
-    [1,             3,                         8,            2,          'iter'],
-    [4,             4,                         16,           4,          'iter'],
-    [4,             2,                         12,           8,          'iter'],
-])
+        [1, 1, 4, 10, 'map'],
+        [1, 3, 8, 2, 'map'],
+        [4, 2, 12, 8, 'map'],
+        [4, 4, 16, 4, 'map'],
+        [1, 1, 4, 10, 'iter'],
+        [1, 3, 8, 2, 'iter'],
+        [4, 4, 16, 4, 'iter'],
+        [4, 2, 12, 8, 'iter'],
+    ]
+)
 def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expected_max_steps, ds_type):
 
     hparams = Namespace(
@@ -100,7 +102,7 @@ def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expe
     tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, tokenizer=tokenizer)
@@ -110,22 +112,20 @@ def test_fix_max_steps_cpu(max_epochs, accumulate_grad_batches, batch_size, expe
     assert model.computed_max_steps == expected_max_steps
 
 
-
 # Test if max_steps fix works correctly
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize(
     ["max_epochs", "accumulate_grad_batches", "batch_size", "distributed_backend", "expected_max_steps"], [
-
-    [1,             1,                         4,            'dp',                  10],
-    [1,             3,                         8,            'dp',                  2],
-    [4,             2,                         12,           'dp',                  8],
-    [4,             4,                         16,           'dp',                  4],
-
-    [1,             1,                         4,            'ddp',                  5],
-    [1,             3,                         8,            'ddp',                  1],
-    [4,             2,                         12,           'ddp',                  4],
-    [4,             4,                         16,           'ddp',                  4],
-])
+        [1, 1, 4, 'dp', 10],
+        [1, 3, 8, 'dp', 2],
+        [4, 2, 12, 'dp', 8],
+        [4, 4, 16, 'dp', 4],
+        [1, 1, 4, 'ddp', 5],
+        [1, 3, 8, 'ddp', 1],
+        [4, 2, 12, 'ddp', 4],
+        [4, 4, 16, 'ddp', 4],
+    ]
+)
 def test_fix_max_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, distributed_backend, expected_max_steps):
 
     time.sleep(5)
@@ -155,12 +155,12 @@ def test_fix_max_steps_gpu(max_epochs, accumulate_grad_batches, batch_size, dist
     )
 
     # instantiate PL model
-    model = SimpleTransformerLikeModel(hparams)    
+    model = SimpleTransformerLikeModel(hparams)
 
     # Datasets
     datamodule = ExampleDataModule(hparams, tokenizer=tokenizer)
 
     # Train!
     trainer.fit(model, datamodule=datamodule)
-    
-    assert model.computed_max_steps == expected_max_steps
\ No newline at end of file
+
+    assert model.computed_max_steps == expected_max_steps
diff --git a/tests/schedulers/test_constant_scheduler.py b/tests/schedulers/test_constant_scheduler.py
index f9870ca..2486e77 100644
--- a/tests/schedulers/test_constant_scheduler.py
+++ b/tests/schedulers/test_constant_scheduler.py
@@ -19,22 +19,17 @@ def __init__(self, hparams):
         self.lrs = []
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = ConstantScheduler(optimizer,
-                                      last_epoch=self.hparams.last_epoch)
+        scheduler = ConstantScheduler(optimizer, last_epoch=self.hparams.last_epoch)
 
         return {
             'optimizer': optimizer,
@@ -56,13 +51,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -70,14 +62,11 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["num_training_steps", "last_epoch", "expected_lrs"], [
-        [20, -1, [1.0] * 21],
-        [100, -1, [1.0] * 101],
-        [1, -1, [1.0, 1.0]]
-])
+    ["num_training_steps", "last_epoch", "expected_lrs"],
+    [[20, -1, [1.0] * 21], [100, -1, [1.0] * 101], [1, -1, [1.0, 1.0]]]
+)
 def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -105,7 +94,7 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -113,6 +102,4 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert expected_lrs == model.lrs, (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert expected_lrs == model.lrs, (f"{expected_lrs} vs {model.lrs}")
diff --git a/tests/schedulers/test_constant_scheduler_with_warmup.py b/tests/schedulers/test_constant_scheduler_with_warmup.py
index 30ee4c7..20d76fa 100644
--- a/tests/schedulers/test_constant_scheduler_with_warmup.py
+++ b/tests/schedulers/test_constant_scheduler_with_warmup.py
@@ -19,23 +19,17 @@ def __init__(self, hparams):
         self.lrs = []
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = ConstantSchedulerWithWarmup(optimizer,
-                                                num_warmup_steps=5,
-                                                last_epoch=self.hparams.last_epoch)
+        scheduler = ConstantSchedulerWithWarmup(optimizer, num_warmup_steps=5, last_epoch=self.hparams.last_epoch)
 
         return {
             'optimizer': optimizer,
@@ -57,13 +51,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -71,12 +62,10 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
-    ["num_training_steps", "last_epoch", "expected_lrs"], [
-        [20, -1, [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] + [1.0] * 15]
-])
+    ["num_training_steps", "last_epoch", "expected_lrs"], [[20, -1, [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] + [1.0] * 15]]
+)
 def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -104,7 +93,7 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -112,6 +101,4 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert expected_lrs == model.lrs, (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert expected_lrs == model.lrs, (f"{expected_lrs} vs {model.lrs}")
diff --git a/tests/schedulers/test_cosine_scheduler_with_warmup.py b/tests/schedulers/test_cosine_scheduler_with_warmup.py
index c5482be..17cc839 100644
--- a/tests/schedulers/test_cosine_scheduler_with_warmup.py
+++ b/tests/schedulers/test_cosine_scheduler_with_warmup.py
@@ -19,25 +19,23 @@ def __init__(self, hparams):
         self.lrs = []
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = CosineSchedulerWithWarmup(optimizer,
-                                              num_warmup_steps=5,
-                                              num_training_steps=self.hparams.max_steps,
-                                              num_cycles=0.5,
-                                              last_epoch=self.hparams.last_epoch)
+        scheduler = CosineSchedulerWithWarmup(
+            optimizer,
+            num_warmup_steps=5,
+            num_training_steps=self.hparams.max_steps,
+            num_cycles=0.5,
+            last_epoch=self.hparams.last_epoch
+        )
 
         return {
             'optimizer': optimizer,
@@ -59,13 +57,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -73,12 +68,18 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
     ["num_training_steps", "last_epoch", "expected_lrs"], [
-        [20, -1, [0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.989074, 0.956773, 0.904508, 0.834565, 0.75, 0.654508, 0.552264, 0.447735, 0.345491, 0.25, 0.165435, 0.095491, 0.043227, 0.010926, 0.0]]
-])
+        [
+            20, -1,
+            [
+                0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.989074, 0.956773, 0.904508, 0.834565, 0.75, 0.654508, 0.552264,
+                0.447735, 0.345491, 0.25, 0.165435, 0.095491, 0.043227, 0.010926, 0.0
+            ]
+        ]
+    ]
+)
 def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -106,7 +107,7 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -114,6 +115,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert torch.tensor(expected_lrs, dtype=torch.float32).allclose(torch.tensor(model.lrs, dtype=torch.float32), rtol=1e-04), (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert torch.tensor(expected_lrs,
+                        dtype=torch.float32).allclose(torch.tensor(model.lrs, dtype=torch.float32),
+                                                      rtol=1e-04), (f"{expected_lrs} vs {model.lrs}")
diff --git a/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py b/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
index 45e3ecc..532ea3a 100644
--- a/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
+++ b/tests/schedulers/test_cosine_scheduler_with_warmup_and_hard_reset.py
@@ -19,25 +19,23 @@ def __init__(self, hparams):
         self.lrs = []
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = CosineSchedulerWithWarmupAndHardRestart(optimizer,
-                                                            num_warmup_steps=5,
-                                                            num_training_steps=self.hparams.max_steps,
-                                                            num_cycles=2.0,
-                                                            last_epoch=self.hparams.last_epoch)
+        scheduler = CosineSchedulerWithWarmupAndHardRestart(
+            optimizer,
+            num_warmup_steps=5,
+            num_training_steps=self.hparams.max_steps,
+            num_cycles=2.0,
+            last_epoch=self.hparams.last_epoch
+        )
 
         return {
             'optimizer': optimizer,
@@ -59,13 +57,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -73,12 +68,18 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
     ["num_training_steps", "last_epoch", "expected_lrs"], [
-        [20, -1, [0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.956772, 0.834565, 0.654508, 0.447736, 0.25, 0.095491, 0.010926, 0.989074, 0.904508, 0.75, 0.552264, 0.345491, 0.165435, 0.043227, 0.0]]
-])
+        [
+            20, -1,
+            [
+                0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.956772, 0.834565, 0.654508, 0.447736, 0.25, 0.095491, 0.010926,
+                0.989074, 0.904508, 0.75, 0.552264, 0.345491, 0.165435, 0.043227, 0.0
+            ]
+        ]
+    ]
+)
 def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -106,7 +107,7 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -114,6 +115,6 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert torch.tensor(expected_lrs, dtype=torch.float32).allclose(torch.tensor(model.lrs, dtype=torch.float32), rtol=1e-04), (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert torch.tensor(expected_lrs,
+                        dtype=torch.float32).allclose(torch.tensor(model.lrs, dtype=torch.float32),
+                                                      rtol=1e-04), (f"{expected_lrs} vs {model.lrs}")
diff --git a/tests/schedulers/test_linear_scheduler.py b/tests/schedulers/test_linear_scheduler.py
index 2989db5..ab623ac 100644
--- a/tests/schedulers/test_linear_scheduler.py
+++ b/tests/schedulers/test_linear_scheduler.py
@@ -18,23 +18,19 @@ def __init__(self, hparams):
         self.model = torch.nn.Linear(10, 10)
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = LinearScheduler(optimizer,
-                                    num_training_steps=self.hparams.max_steps,
-                                    last_epoch=self.hparams.last_epoch)
+        scheduler = LinearScheduler(
+            optimizer, num_training_steps=self.hparams.max_steps, last_epoch=self.hparams.last_epoch
+        )
 
         return {
             'optimizer': optimizer,
@@ -56,13 +52,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -70,12 +63,18 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
     ["num_training_steps", "last_epoch", "expected_lrs"], [
-        [20, -1, [1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1, 0.05, 0.0]]
-])
+        [
+            20, -1,
+            [
+                1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1,
+                0.05, 0.0
+            ]
+        ]
+    ]
+)
 def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -103,7 +102,7 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -111,6 +110,4 @@ def test_datamodule_cpu(num_training_steps, last_epoch, expected_lrs):
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert expected_lrs == model.lrs, (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert expected_lrs == model.lrs, (f"{expected_lrs} vs {model.lrs}")
diff --git a/tests/schedulers/test_linear_scheduler_with_warmup.py b/tests/schedulers/test_linear_scheduler_with_warmup.py
index 08c527e..79ceb18 100644
--- a/tests/schedulers/test_linear_scheduler_with_warmup.py
+++ b/tests/schedulers/test_linear_scheduler_with_warmup.py
@@ -18,24 +18,22 @@ def __init__(self, hparams):
         self.model = torch.nn.Linear(10, 10)
 
     def training_step(self, batch, *args):
-        return {
-            'loss': self(batch['data']).sum(),
-            'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']
-        }
-    
+        return {'loss': self(batch['data']).sum(), 'lr': self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']}
+
     def training_epoch_end(self, outputs, *args, **kwargs):
         self.lrs = [o['lr'] for o in outputs] + [self.trainer.optimizers[0].__dict__['param_groups'][0]['lr']]
 
     def configure_optimizers(self):
         # Define adam optimizer
-        optimizer = AdamW(self.model.parameters(),
-                          lr=self.hparams.learning_rate)
+        optimizer = AdamW(self.model.parameters(), lr=self.hparams.learning_rate)
 
         # init scheduler after optional fp16 to get rid of strange warning about optimizer and scheduler steps order
-        scheduler = LinearSchedulerWithWarmup(optimizer,
-                                              num_warmup_steps=self.hparams.num_warmup_steps,
-                                              num_training_steps=self.hparams.max_steps,
-                                              last_epoch=self.hparams.last_epoch)
+        scheduler = LinearSchedulerWithWarmup(
+            optimizer,
+            num_warmup_steps=self.hparams.num_warmup_steps,
+            num_training_steps=self.hparams.max_steps,
+            last_epoch=self.hparams.last_epoch
+        )
 
         return {
             'optimizer': optimizer,
@@ -57,13 +55,10 @@ def __iter__(self):
             yield x
 
     def preprocess_line(self, line: list) -> list:
-        res = {
-            'data': [float(line)] * 10
-        }
+        res = {'data': [float(line)] * 10}
         return res
 
 
-
 class FakeDataModule(SuperDataModule):
 
     def __init__(self, *args, **kwargs):
@@ -71,12 +66,15 @@ def __init__(self, *args, **kwargs):
         self.train_adapter = FakeAdapter(self.hparams)
 
 
-
 # Test iter dataset work correctly
 @pytest.mark.parametrize(
     ["num_warmup_steps", "num_training_steps", "last_epoch", "expected_lrs"], [
-        [10, 20, -1, [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0]]
-])
+        [
+            10, 20, -1,
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0]
+        ]
+    ]
+)
 def test_datamodule_cpu(num_warmup_steps, num_training_steps, last_epoch, expected_lrs):
 
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -105,7 +103,7 @@ def test_datamodule_cpu(num_warmup_steps, num_training_steps, last_epoch, expect
     )
 
     # instantiate PL model
-    model = FakeModel(hparams)    
+    model = FakeModel(hparams)
 
     # Datasets
     datamodule = FakeDataModule(hparams)
@@ -113,6 +111,4 @@ def test_datamodule_cpu(num_warmup_steps, num_training_steps, last_epoch, expect
     # Fit
     trainer.fit(model, datamodule=datamodule)
 
-    assert expected_lrs == model.lrs, (
-        f"{expected_lrs} vs {model.lrs}"
-    )
+    assert expected_lrs == model.lrs, (f"{expected_lrs} vs {model.lrs}")