From aaf22e93d95d13fddd3b25f8543929fb1ab077bc Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Mon, 28 Mar 2022 16:15:03 +0200
Subject: [PATCH 01/15] Avoid accessing .dataset of a dataloader

---
 src/transformers/trainer.py | 61 +++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 1bf6fde9fc5e..da7dc5e9d571 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -938,10 +938,12 @@ def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optim
     def num_examples(self, dataloader: DataLoader) -> int:
         """
         Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset.
-
-        Will raise an exception if the underlying dataset does not implement method `__len__`
+        When dataloader.dataset does not exist or has no length, estimates as best it can
         """
-        return len(dataloader.dataset)
+        try:
+            return len(dataloader.dataset)
+        except (NameError, AttributeError, TypeError):  # no dataset or length, estimate by length of dataloader
+            return len(dataloader) * self.args.per_device_train_batch_size
 
     def _hp_search_setup(self, trial: Union["optuna.Trial", Dict[str, Any]]):
         """HP search setup code"""
@@ -1198,39 +1200,45 @@ def train(
                 self._move_model_to_device(self.model, args.device)
             self.model_wrapped = self.model
 
-        # Keeping track whether we can can len() on the dataset or not
-        train_dataset_is_sized = has_length(self.train_dataset)
-
         # Data loader and number of training steps
         train_dataloader = self.get_train_dataloader()
 
+        # Keeping track whether we can can len() on the dataset or not
+
+
+
         # Setting up training control variables:
         # number of training epochs: num_train_epochs
         # number of training steps per epoch: num_update_steps_per_epoch
         # total number of training steps to execute: max_steps
         total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * args.world_size
-        if train_dataset_is_sized:
-            num_update_steps_per_epoch = len(train_dataloader) // args.gradient_accumulation_steps
+
+        len_dataloader = None
+        try:
+            len_dataloader = len(train_dataloader)
+        except (NameError, TypeError):  # Not sure when this happens, don't all dataloaders have len() ?
+            # see __init__. max_steps is set when the dataset has no __len__
+            max_steps = args.max_steps
+            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+            num_train_epochs = sys.maxsize
+            num_update_steps_per_epoch = max_steps
+            num_train_samples = args.max_steps * total_train_batch_size
+        else:
+            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
             num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
             if args.max_steps > 0:
                 max_steps = args.max_steps
                 num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                     args.max_steps % num_update_steps_per_epoch > 0
                 )
-                # May be slightly incorrect if the last batch in the training datalaoder has a smaller size but it's
+                # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
                 # the best we can do.
                 num_train_samples = args.max_steps * total_train_batch_size
             else:
                 max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                 num_train_epochs = math.ceil(args.num_train_epochs)
-                num_train_samples = len(self.train_dataset) * args.num_train_epochs
-        else:
-            # see __init__. max_steps is set when the dataset has no __len__
-            max_steps = args.max_steps
-            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
-            num_train_epochs = sys.maxsize
-            num_update_steps_per_epoch = max_steps
-            num_train_samples = args.max_steps * total_train_batch_size
+                num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
+
 
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
@@ -1281,12 +1289,8 @@ def train(
         # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.
 
         # Train!
-        num_examples = (
-            self.num_examples(train_dataloader) if train_dataset_is_sized else total_train_batch_size * args.max_steps
-        )
-
         logger.info("***** Running training *****")
-        logger.info(f"  Num examples = {num_examples}")
+        logger.info(f"  Num examples = {num_train_samples}")
         logger.info(f"  Num Epochs = {num_train_epochs}")
         logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
         logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
@@ -1370,7 +1374,7 @@ def train(
         for epoch in range(epochs_trained, num_train_epochs):
             if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                 train_dataloader.sampler.set_epoch(epoch)
-            elif isinstance(train_dataloader.dataset, IterableDatasetShard):
+            elif hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDatasetShard):
                 train_dataloader.dataset.set_epoch(epoch)
 
             if is_torch_tpu_available():
@@ -1384,7 +1388,7 @@ def train(
                 self._past = None
 
             steps_in_epoch = (
-                len(epoch_iterator) if train_dataset_is_sized else args.max_steps * args.gradient_accumulation_steps
+                len(epoch_iterator) if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps
             )
             self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
 
@@ -2420,7 +2424,6 @@ def evaluation_loop(
 
         self.callback_handler.eval_dataloader = dataloader
         # Do this before wrapping.
-        eval_dataset = dataloader.dataset
 
         if is_torch_tpu_available():
             dataloader = pl.ParallelLoader(dataloader, [args.device]).per_device_loader(args.device)
@@ -2505,14 +2508,14 @@ def evaluation_loop(
             all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
 
         # Number of samples
-        if has_length(eval_dataset):
-            num_samples = len(eval_dataset)
+        eval_dataset = getattr(dataloader, 'dataset',None)
+
         # The instance check is weird and does not actually check for the type, but whether the dataset has the right
         # methods. Therefore we need to make sure it also has the attribute.
-        elif isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
+        if isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
             num_samples = eval_dataset.num_examples
         else:
-            num_samples = observed_num_examples
+            num_samples = self.num_examples(dataloader)
 
         # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
         # samplers has been rounded to a multiple of batch_size, so we truncate.

From c87e3f71b440f02dec41ad6032770aea173cbbad Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Mon, 28 Mar 2022 16:17:45 +0200
Subject: [PATCH 02/15] style

---
 src/transformers/trainer.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index da7dc5e9d571..78b758209999 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -937,8 +937,8 @@ def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optim
 
     def num_examples(self, dataloader: DataLoader) -> int:
         """
-        Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset.
-        When dataloader.dataset does not exist or has no length, estimates as best it can
+        Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset. When
+        dataloader.dataset does not exist or has no length, estimates as best it can
         """
         try:
             return len(dataloader.dataset)
@@ -1205,8 +1205,6 @@ def train(
 
         # Keeping track whether we can can len() on the dataset or not
 
-
-
         # Setting up training control variables:
         # number of training epochs: num_train_epochs
         # number of training steps per epoch: num_update_steps_per_epoch
@@ -1239,7 +1237,6 @@ def train(
                 num_train_epochs = math.ceil(args.num_train_epochs)
                 num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
 
-
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
                 # nn.DataParallel(model) replicates the model, creating new variables and module
@@ -1388,7 +1385,9 @@ def train(
                 self._past = None
 
             steps_in_epoch = (
-                len(epoch_iterator) if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps
+                len(epoch_iterator)
+                if len_dataloader is not None
+                else args.max_steps * args.gradient_accumulation_steps
             )
             self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
 
@@ -2508,7 +2507,7 @@ def evaluation_loop(
             all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
 
         # Number of samples
-        eval_dataset = getattr(dataloader, 'dataset',None)
+        eval_dataset = getattr(dataloader, "dataset", None)
 
         # The instance check is weird and does not actually check for the type, but whether the dataset has the right
         # methods. Therefore we need to make sure it also has the attribute.

From 41598a67dbc6ae302caea1f2c8617087c0d5298e Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Mon, 28 Mar 2022 16:45:22 +0200
Subject: [PATCH 03/15] fix

---
 src/transformers/trainer.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 78b758209999..4aad95740750 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1203,8 +1203,6 @@ def train(
         # Data loader and number of training steps
         train_dataloader = self.get_train_dataloader()
 
-        # Keeping track whether we can can len() on the dataset or not
-
         # Setting up training control variables:
         # number of training epochs: num_train_epochs
         # number of training steps per epoch: num_update_steps_per_epoch
@@ -1214,7 +1212,7 @@ def train(
         len_dataloader = None
         try:
             len_dataloader = len(train_dataloader)
-        except (NameError, TypeError):  # Not sure when this happens, don't all dataloaders have len() ?
+        except (NameError, TypeError):  # Default dataloader calls len(dataset), which may not exist
             # see __init__. max_steps is set when the dataset has no __len__
             max_steps = args.max_steps
             # Setting a very large number of epochs so we go as many times as necessary over the iterator.
@@ -2423,6 +2421,7 @@ def evaluation_loop(
 
         self.callback_handler.eval_dataloader = dataloader
         # Do this before wrapping.
+        eval_dataset = getattr(dataloader, "dataset", None)
 
         if is_torch_tpu_available():
             dataloader = pl.ParallelLoader(dataloader, [args.device]).per_device_loader(args.device)
@@ -2507,14 +2506,15 @@ def evaluation_loop(
             all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
 
         # Number of samples
-        eval_dataset = getattr(dataloader, "dataset", None)
-
         # The instance check is weird and does not actually check for the type, but whether the dataset has the right
         # methods. Therefore we need to make sure it also has the attribute.
         if isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
             num_samples = eval_dataset.num_examples
         else:
-            num_samples = self.num_examples(dataloader)
+            try:
+                num_samples = self.num_examples(dataloader)
+            except TypeError:  # both len(dataloader.dataset) and len(dataloader) fail
+                num_samples = observed_num_examples
 
         # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
         # samplers has been rounded to a multiple of batch_size, so we truncate.
@@ -2901,8 +2901,11 @@ def prediction_loop(
         """
         args = self.args
 
-        if not has_length(dataloader.dataset):
-            raise ValueError("dataset must implement __len__")
+        try:
+            num_examples = self.num_examples(dataloader)
+        except TypeError:
+            raise ValueError("dataloader or it's dataset must implement __len__")
+
         prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only
 
         # if eval is called w/o train init deepspeed here
@@ -2931,7 +2934,6 @@ def prediction_loop(
                 model = model.to(dtype=torch.bfloat16, device=args.device)
 
         batch_size = dataloader.batch_size
-        num_examples = self.num_examples(dataloader)
         logger.info(f"***** Running {description} *****")
         logger.info(f"  Num examples = {num_examples}")
         logger.info(f"  Batch size = {batch_size}")

From 00b82f64e42b2f9a54d961e40b72a2342f0da004 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Mon, 28 Mar 2022 17:08:10 +0200
Subject: [PATCH 04/15] cleaning up, reverting some misunderstandings

---
 src/transformers/trainer.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 4aad95740750..ea8367d440c8 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1218,10 +1218,12 @@ def train(
             # Setting a very large number of epochs so we go as many times as necessary over the iterator.
             num_train_epochs = sys.maxsize
             num_update_steps_per_epoch = max_steps
+            num_examples = total_train_batch_size * args.max_steps
             num_train_samples = args.max_steps * total_train_batch_size
         else:
             num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
             num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
+            num_examples = self.num_examples(train_dataloader)
             if args.max_steps > 0:
                 max_steps = args.max_steps
                 num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
@@ -1235,6 +1237,7 @@ def train(
                 num_train_epochs = math.ceil(args.num_train_epochs)
                 num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
 
+
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
                 # nn.DataParallel(model) replicates the model, creating new variables and module
@@ -1285,7 +1288,7 @@ def train(
 
         # Train!
         logger.info("***** Running training *****")
-        logger.info(f"  Num examples = {num_train_samples}")
+        logger.info(f"  Num examples = {num_examples}")
         logger.info(f"  Num Epochs = {num_train_epochs}")
         logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
         logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
@@ -2506,9 +2509,11 @@ def evaluation_loop(
             all_labels = labels if all_labels is None else nested_concat(all_labels, labels, padding_index=-100)
 
         # Number of samples
+        if has_length(eval_dataset):
+            num_samples = len(eval_dataset)
         # The instance check is weird and does not actually check for the type, but whether the dataset has the right
         # methods. Therefore we need to make sure it also has the attribute.
-        if isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
+        elif isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
             num_samples = eval_dataset.num_examples
         else:
             try:
@@ -2904,7 +2909,7 @@ def prediction_loop(
         try:
             num_examples = self.num_examples(dataloader)
         except TypeError:
-            raise ValueError("dataloader or it's dataset must implement __len__")
+            raise ValueError("dataloader or its dataset must implement __len__")
 
         prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only
 

From 4ee0a57662ed024028da38b5709376a22825413f Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Mon, 28 Mar 2022 17:09:57 +0200
Subject: [PATCH 05/15] black

---
 src/transformers/trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index ea8367d440c8..9049069b6472 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1237,7 +1237,6 @@ def train(
                 num_train_epochs = math.ceil(args.num_train_epochs)
                 num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
 
-
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
                 # nn.DataParallel(model) replicates the model, creating new variables and module

From 815960cd119c191049e4033ba94d3f12b7a486db Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 12:05:27 +0200
Subject: [PATCH 06/15] add train_dataset argument to get_train_dataloader, and
 fix other instances of length checks

---
 src/transformers/trainer.py        | 46 ++++++++++++++++++------------
 src/transformers/utils/notebook.py |  8 ++++--
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 9049069b6472..adae1b238628 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -584,8 +584,10 @@ def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optio
         else:
             return dataset.remove_columns(ignored_columns)
 
-    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
-        if not has_length(self.train_dataset):
+    def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]:
+        train_dataset = train_dataset if train_dataset is not None else self.train_dataset
+
+        if not has_length(train_dataset):
             return None
 
         generator = None
@@ -604,10 +606,10 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
 
         # Build the sampler.
         if self.args.group_by_length:
-            if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
+            if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
                 lengths = (
-                    self.train_dataset[self.args.length_column_name]
-                    if self.args.length_column_name in self.train_dataset.column_names
+                    train_dataset[self.args.length_column_name]
+                    if self.args.length_column_name in train_dataset.column_names
                     else None
                 )
             else:
@@ -616,7 +618,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
             if self.args.world_size <= 1:
                 return LengthGroupedSampler(
                     self.args.train_batch_size * self.args.gradient_accumulation_steps,
-                    dataset=self.train_dataset,
+                    dataset=train_dataset,
                     lengths=lengths,
                     model_input_name=model_input_name,
                     generator=generator,
@@ -624,7 +626,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
             else:
                 return DistributedLengthGroupedSampler(
                     self.args.train_batch_size * self.args.gradient_accumulation_steps,
-                    dataset=self.train_dataset,
+                    dataset=train_dataset,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
                     lengths=lengths,
@@ -635,15 +637,15 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
         else:
             if self.args.world_size <= 1:
                 if _is_torch_generator_available:
-                    return RandomSampler(self.train_dataset, generator=generator)
-                return RandomSampler(self.train_dataset)
+                    return RandomSampler(train_dataset, generator=generator)
+                return RandomSampler(train_dataset)
             elif (
                 self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL]
                 and not self.args.dataloader_drop_last
             ):
                 # Use a loop for TPUs when drop_last is False to have all batches have the same size.
                 return DistributedSamplerWithLoop(
-                    self.train_dataset,
+                    train_dataset,
                     batch_size=self.args.per_device_train_batch_size,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
@@ -651,25 +653,31 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
                 )
             else:
                 return DistributedSampler(
-                    self.train_dataset,
+                    train_dataset,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
                     seed=seed,
                 )
 
-    def get_train_dataloader(self) -> DataLoader:
+    def get_train_dataloader(self, train_dataset: Optional[Dataset] = None) -> DataLoader:
         """
         Returns the training [`~torch.utils.data.DataLoader`].
 
-        Will use no sampler if `self.train_dataset` does not implement `__len__`, a random sampler (adapted to
-        distributed training if necessary) otherwise.
+        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
+        training if necessary) otherwise.
 
         Subclass and override this method if you want to inject some custom behavior.
+
+        Args:
+            train_dataset (`torch.utils.data.Dataset`, *optional*):
+                If provided, will override `self.train_dataset`. If it is an `datasets.Dataset`, columns not accepted
+                by the `model.forward()` method are automatically removed.
         """
-        if self.train_dataset is None:
+
+        if train_dataset is None and self.train_dataset is None:
             raise ValueError("Trainer: training requires a train_dataset.")
+        train_dataset = train_dataset if train_dataset is not None else self.train_dataset
 
-        train_dataset = self.train_dataset
         if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
             train_dataset = self._remove_unused_columns(train_dataset, description="training")
 
@@ -691,7 +699,7 @@ def get_train_dataloader(self) -> DataLoader:
                 pin_memory=self.args.dataloader_pin_memory,
             )
 
-        train_sampler = self._get_train_sampler()
+        train_sampler = self._get_train_sampler(train_dataset)
 
         return DataLoader(
             train_dataset,
@@ -2413,9 +2421,9 @@ def evaluation_loop(
         batch_size = dataloader.batch_size
 
         logger.info(f"***** Running {description} *****")
-        if has_length(dataloader.dataset):
+        try:
             logger.info(f"  Num examples = {self.num_examples(dataloader)}")
-        else:
+        except TypeError:
             logger.info("  Num examples: Unknown")
         logger.info(f"  Batch size = {batch_size}")
 
diff --git a/src/transformers/utils/notebook.py b/src/transformers/utils/notebook.py
index 779446f5f104..e0b259b53445 100644
--- a/src/transformers/utils/notebook.py
+++ b/src/transformers/utils/notebook.py
@@ -294,13 +294,15 @@ def on_step_end(self, args, state, control, **kwargs):
         self._force_next_update = False
 
     def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
-        if not isinstance(eval_dataloader.dataset, collections.abc.Sized):
+        try:
+            len_dataloader = len(eval_dataloader)
+        except (NameError, TypeError, AttributeError):
             return
         if self.prediction_bar is None:
             if self.training_tracker is not None:
-                self.prediction_bar = self.training_tracker.add_child(len(eval_dataloader))
+                self.prediction_bar = self.training_tracker.add_child(len_dataloader)
             else:
-                self.prediction_bar = NotebookProgressBar(len(eval_dataloader))
+                self.prediction_bar = NotebookProgressBar(len_dataloader)
             self.prediction_bar.update(1)
         else:
             self.prediction_bar.update(self.prediction_bar.value + 1)

From 4c02f8c8e0b358837751d3dad37921936d0b3359 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 12:16:16 +0200
Subject: [PATCH 07/15] flake8

---
 src/transformers/utils/notebook.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/transformers/utils/notebook.py b/src/transformers/utils/notebook.py
index e0b259b53445..1e257b4db67a 100644
--- a/src/transformers/utils/notebook.py
+++ b/src/transformers/utils/notebook.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import collections
 import re
 import time
 from typing import Optional

From 87f45ec81f0a555da7fa779e7f43eb29d34792f2 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 17:54:56 +0200
Subject: [PATCH 08/15] address comments

---
 src/transformers/trainer.py        | 70 +++++++++++++-----------------
 src/transformers/utils/notebook.py | 10 ++---
 2 files changed, 35 insertions(+), 45 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index adae1b238628..b87e2dd89cac 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -584,10 +584,9 @@ def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optio
         else:
             return dataset.remove_columns(ignored_columns)
 
-    def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]:
-        train_dataset = train_dataset if train_dataset is not None else self.train_dataset
+    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
 
-        if not has_length(train_dataset):
+        if self.train_dataset is None or not has_length(self.train_dataset):
             return None
 
         generator = None
@@ -606,10 +605,10 @@ def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optiona
 
         # Build the sampler.
         if self.args.group_by_length:
-            if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
+            if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
                 lengths = (
-                    train_dataset[self.args.length_column_name]
-                    if self.args.length_column_name in train_dataset.column_names
+                    self.train_dataset[self.args.length_column_name]
+                    if self.args.length_column_name in self.train_dataset.column_names
                     else None
                 )
             else:
@@ -618,7 +617,7 @@ def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optiona
             if self.args.world_size <= 1:
                 return LengthGroupedSampler(
                     self.args.train_batch_size * self.args.gradient_accumulation_steps,
-                    dataset=train_dataset,
+                    dataset=self.train_dataset,
                     lengths=lengths,
                     model_input_name=model_input_name,
                     generator=generator,
@@ -626,7 +625,7 @@ def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optiona
             else:
                 return DistributedLengthGroupedSampler(
                     self.args.train_batch_size * self.args.gradient_accumulation_steps,
-                    dataset=train_dataset,
+                    dataset=self.train_dataset,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
                     lengths=lengths,
@@ -637,15 +636,15 @@ def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optiona
         else:
             if self.args.world_size <= 1:
                 if _is_torch_generator_available:
-                    return RandomSampler(train_dataset, generator=generator)
-                return RandomSampler(train_dataset)
+                    return RandomSampler(self.train_dataset, generator=generator)
+                return RandomSampler(self.train_dataset)
             elif (
                 self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL]
                 and not self.args.dataloader_drop_last
             ):
                 # Use a loop for TPUs when drop_last is False to have all batches have the same size.
                 return DistributedSamplerWithLoop(
-                    train_dataset,
+                    self.train_dataset,
                     batch_size=self.args.per_device_train_batch_size,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
@@ -653,13 +652,13 @@ def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optiona
                 )
             else:
                 return DistributedSampler(
-                    train_dataset,
+                    self.train_dataset,
                     num_replicas=self.args.world_size,
                     rank=self.args.process_index,
                     seed=seed,
                 )
 
-    def get_train_dataloader(self, train_dataset: Optional[Dataset] = None) -> DataLoader:
+    def get_train_dataloader(self) -> DataLoader:
         """
         Returns the training [`~torch.utils.data.DataLoader`].
 
@@ -667,18 +666,13 @@ def get_train_dataloader(self, train_dataset: Optional[Dataset] = None) -> DataL
         training if necessary) otherwise.
 
         Subclass and override this method if you want to inject some custom behavior.
-
-        Args:
-            train_dataset (`torch.utils.data.Dataset`, *optional*):
-                If provided, will override `self.train_dataset`. If it is an `datasets.Dataset`, columns not accepted
-                by the `model.forward()` method are automatically removed.
         """
 
-        if train_dataset is None and self.train_dataset is None:
+        if self.train_dataset is None:
             raise ValueError("Trainer: training requires a train_dataset.")
-        train_dataset = train_dataset if train_dataset is not None else self.train_dataset
 
-        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
+        train_dataset = self.train_dataset
+        if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
             train_dataset = self._remove_unused_columns(train_dataset, description="training")
 
         if isinstance(train_dataset, torch.utils.data.IterableDataset):
@@ -1218,17 +1212,8 @@ def train(
         total_train_batch_size = args.train_batch_size * args.gradient_accumulation_steps * args.world_size
 
         len_dataloader = None
-        try:
+        if has_length(train_dataloader):
             len_dataloader = len(train_dataloader)
-        except (NameError, TypeError):  # Default dataloader calls len(dataset), which may not exist
-            # see __init__. max_steps is set when the dataset has no __len__
-            max_steps = args.max_steps
-            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
-            num_train_epochs = sys.maxsize
-            num_update_steps_per_epoch = max_steps
-            num_examples = total_train_batch_size * args.max_steps
-            num_train_samples = args.max_steps * total_train_batch_size
-        else:
             num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
             num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
             num_examples = self.num_examples(train_dataloader)
@@ -1244,6 +1229,14 @@ def train(
                 max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                 num_train_epochs = math.ceil(args.num_train_epochs)
                 num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
+        else:
+            # see __init__. max_steps is set when the dataset has no __len__
+            max_steps = args.max_steps
+            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
+            num_train_epochs = sys.maxsize
+            num_update_steps_per_epoch = max_steps
+            num_examples = total_train_batch_size * args.max_steps
+            num_train_samples = args.max_steps * total_train_batch_size
 
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
@@ -2421,9 +2414,9 @@ def evaluation_loop(
         batch_size = dataloader.batch_size
 
         logger.info(f"***** Running {description} *****")
-        try:
+        if has_length(dataloader):
             logger.info(f"  Num examples = {self.num_examples(dataloader)}")
-        except TypeError:
+        else:
             logger.info("  Num examples: Unknown")
         logger.info(f"  Batch size = {batch_size}")
 
@@ -2523,9 +2516,9 @@ def evaluation_loop(
         elif isinstance(eval_dataset, IterableDatasetShard) and hasattr(eval_dataset, "num_examples"):
             num_samples = eval_dataset.num_examples
         else:
-            try:
+            if has_length(dataloader):
                 num_samples = self.num_examples(dataloader)
-            except TypeError:  # both len(dataloader.dataset) and len(dataloader) fail
+            else:  # both len(dataloader.dataset) and len(dataloader) fail
                 num_samples = observed_num_examples
 
         # Number of losses has been rounded to a multiple of batch_size and in a distributed training, the number of
@@ -2913,10 +2906,8 @@ def prediction_loop(
         """
         args = self.args
 
-        try:
-            num_examples = self.num_examples(dataloader)
-        except TypeError:
-            raise ValueError("dataloader or its dataset must implement __len__")
+        if not has_length(dataloader):
+            raise ValueError("dataloader must implement a working __len__")
 
         prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only
 
@@ -2946,6 +2937,7 @@ def prediction_loop(
                 model = model.to(dtype=torch.bfloat16, device=args.device)
 
         batch_size = dataloader.batch_size
+        num_examples = self.num_examples(dataloader)
         logger.info(f"***** Running {description} *****")
         logger.info(f"  Num examples = {num_examples}")
         logger.info(f"  Batch size = {batch_size}")
diff --git a/src/transformers/utils/notebook.py b/src/transformers/utils/notebook.py
index 1e257b4db67a..0ffbdc8deecf 100644
--- a/src/transformers/utils/notebook.py
+++ b/src/transformers/utils/notebook.py
@@ -20,7 +20,7 @@
 import IPython.display as disp
 
 from ..trainer_callback import TrainerCallback
-from ..trainer_utils import IntervalStrategy
+from ..trainer_utils import IntervalStrategy, has_length
 
 
 def format_time(t):
@@ -293,15 +293,13 @@ def on_step_end(self, args, state, control, **kwargs):
         self._force_next_update = False
 
     def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
-        try:
-            len_dataloader = len(eval_dataloader)
-        except (NameError, TypeError, AttributeError):
+        if not has_length(eval_dataloader):
             return
         if self.prediction_bar is None:
             if self.training_tracker is not None:
-                self.prediction_bar = self.training_tracker.add_child(len_dataloader)
+                self.prediction_bar = self.training_tracker.add_child(len(eval_dataloader))
             else:
-                self.prediction_bar = NotebookProgressBar(len_dataloader)
+                self.prediction_bar = NotebookProgressBar(len(eval_dataloader))
             self.prediction_bar.update(1)
         else:
             self.prediction_bar.update(self.prediction_bar.value + 1)

From 761ab98160510f52efbd02c0300525d95d6396e9 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 18:14:27 +0200
Subject: [PATCH 09/15] fix bug

---
 src/transformers/trainer.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index b87e2dd89cac..82d6494b361a 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -585,7 +585,6 @@ def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optio
             return dataset.remove_columns(ignored_columns)
 
     def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
-
         if self.train_dataset is None or not has_length(self.train_dataset):
             return None
 
@@ -693,7 +692,7 @@ def get_train_dataloader(self) -> DataLoader:
                 pin_memory=self.args.dataloader_pin_memory,
             )
 
-        train_sampler = self._get_train_sampler(train_dataset)
+        train_sampler = self._get_train_sampler()
 
         return DataLoader(
             train_dataset,
@@ -1229,14 +1228,15 @@ def train(
                 max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                 num_train_epochs = math.ceil(args.num_train_epochs)
                 num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
-        else:
-            # see __init__. max_steps is set when the dataset has no __len__
+        elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
             max_steps = args.max_steps
             # Setting a very large number of epochs so we go as many times as necessary over the iterator.
             num_train_epochs = sys.maxsize
             num_update_steps_per_epoch = max_steps
             num_examples = total_train_batch_size * args.max_steps
             num_train_samples = args.max_steps * total_train_batch_size
+        else:
+            raise ValueError(f"args.max_steps must be set if dataloader does not have a length")
 
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:

From eef37c1157db152a6403bc867ef2eb61ede7fbf5 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 18:18:39 +0200
Subject: [PATCH 10/15] cleanup

---
 src/transformers/trainer.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 82d6494b361a..c939a2b68bcb 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -666,12 +666,11 @@ def get_train_dataloader(self) -> DataLoader:
 
         Subclass and override this method if you want to inject some custom behavior.
         """
-
         if self.train_dataset is None:
             raise ValueError("Trainer: training requires a train_dataset.")
 
         train_dataset = self.train_dataset
-        if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
+        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
             train_dataset = self._remove_unused_columns(train_dataset, description="training")
 
         if isinstance(train_dataset, torch.utils.data.IterableDataset):

From a1697241b3c8ec88aac796e8eeec792dcd019715 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 18:57:49 +0200
Subject: [PATCH 11/15] add test

---
 src/transformers/trainer.py          |  6 ++++--
 src/transformers/trainer_callback.py |  2 +-
 tests/trainer/test_trainer.py        | 31 ++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index c939a2b68bcb..993e5248635e 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1235,7 +1235,9 @@ def train(
             num_examples = total_train_batch_size * args.max_steps
             num_train_samples = args.max_steps * total_train_batch_size
         else:
-            raise ValueError(f"args.max_steps must be set if dataloader does not have a length")
+            raise ValueError(
+                f"args.max_steps must be set to a positive value if dataloader does not have a length, was {args.max_steps}"
+            )
 
         if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
             if self.args.n_gpu > 1:
@@ -2410,7 +2412,7 @@ def evaluation_loop(
             elif args.bf16_full_eval:
                 model = model.to(dtype=torch.bfloat16, device=args.device)
 
-        batch_size = dataloader.batch_size
+        batch_size = self.args.per_device_eval_batch_size
 
         logger.info(f"***** Running {description} *****")
         if has_length(dataloader):
diff --git a/src/transformers/trainer_callback.py b/src/transformers/trainer_callback.py
index ec344341bce6..92abe1ed5063 100644
--- a/src/transformers/trainer_callback.py
+++ b/src/transformers/trainer_callback.py
@@ -473,7 +473,7 @@ def on_step_end(self, args, state, control, **kwargs):
             self.current_step = state.global_step
 
     def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
-        if state.is_local_process_zero and has_length(eval_dataloader.dataset):
+        if state.is_local_process_zero and has_length(eval_dataloader):
             if self.prediction_bar is None:
                 self.prediction_bar = tqdm(total=len(eval_dataloader), leave=self.training_bar is None)
             self.prediction_bar.update(1)
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index ec044ea1daa7..06755d024f18 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -169,6 +169,28 @@ def __init__(self, a=0, b=0, double_output=False, **kwargs):
         self.hidden_size = 1
 
 
+class MultiLoader:
+    def __init__(self, loaders):
+        self.loaders = loaders
+
+    def __len__(self):
+        return sum(len(loader) for loader in self.loaders)
+
+    def __iter__(self):
+        for loader in self.loaders:
+            yield from loader
+
+
+class CustomDataloaderTrainer(Trainer):
+    def get_train_dataloader(self):
+        dataloaders = [super(CustomDataloaderTrainer, self).get_train_dataloader() for _ in range(2)]
+        return MultiLoader(dataloaders)
+
+    def get_eval_dataloader(self, eval_dataset):
+        dataloaders = [super(CustomDataloaderTrainer, self).get_eval_dataloader(eval_dataset) for _ in range(2)]
+        return MultiLoader(dataloaders)
+
+
 if is_torch_available():
 
     class SampleIterableDataset(IterableDataset):
@@ -647,6 +669,15 @@ def test_train_and_eval_dataloaders(self):
         new_eval_dataset = RegressionDataset(length=128)
         self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32 * n_gpu))
 
+    # tests that we do not require dataloader to have a .dataset
+    def test_dataloader_without_dataset(self):
+        train_dataset = RegressionDataset(length=128)
+        trainer = CustomDataloaderTrainer(
+            model=RegressionModel(), train_dataset=train_dataset, eval_dataset=train_dataset
+        )
+        trainer.train()
+        trainer.evaluate()
+
     def test_sampler_seed(self):
         # nb: we don't want to inherit from IterableDataset to hit the right code path
         class DummyDataset(torch.utils.data.Dataset):

From 2ca2f3f35e3a916b651085ea740d3d533c33ab51 Mon Sep 17 00:00:00 2001
From: Sander Land <48946947+sanderland@users.noreply.github.com>
Date: Tue, 29 Mar 2022 19:19:01 +0200
Subject: [PATCH 12/15] Update tests/trainer/test_trainer.py

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
---
 tests/trainer/test_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 06755d024f18..8fabaca22b5e 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -669,7 +669,7 @@ def test_train_and_eval_dataloaders(self):
         new_eval_dataset = RegressionDataset(length=128)
         self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // (32 * n_gpu))
 
-    # tests that we do not require dataloader to have a .dataset
+    # tests that we do not require dataloader to have a .dataset attribute
     def test_dataloader_without_dataset(self):
         train_dataset = RegressionDataset(length=128)
         trainer = CustomDataloaderTrainer(

From b36f278d3ffd7fc34904512bebfa25db69d38c57 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 19:20:11 +0200
Subject: [PATCH 13/15] under torch

---
 tests/trainer/test_trainer.py | 44 +++++++++++++++++------------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 06755d024f18..d8db94db4581 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -169,28 +169,6 @@ def __init__(self, a=0, b=0, double_output=False, **kwargs):
         self.hidden_size = 1
 
 
-class MultiLoader:
-    def __init__(self, loaders):
-        self.loaders = loaders
-
-    def __len__(self):
-        return sum(len(loader) for loader in self.loaders)
-
-    def __iter__(self):
-        for loader in self.loaders:
-            yield from loader
-
-
-class CustomDataloaderTrainer(Trainer):
-    def get_train_dataloader(self):
-        dataloaders = [super(CustomDataloaderTrainer, self).get_train_dataloader() for _ in range(2)]
-        return MultiLoader(dataloaders)
-
-    def get_eval_dataloader(self, eval_dataset):
-        dataloaders = [super(CustomDataloaderTrainer, self).get_eval_dataloader(eval_dataset) for _ in range(2)]
-        return MultiLoader(dataloaders)
-
-
 if is_torch_available():
 
     class SampleIterableDataset(IterableDataset):
@@ -211,6 +189,28 @@ def __iter__(self):
                 yield self.dataset[self.current_sample]
                 self.current_sample += 1
 
+
+    class MultiLoader:
+        def __init__(self, loaders):
+            self.loaders = loaders
+
+        def __len__(self):
+            return sum(len(loader) for loader in self.loaders)
+
+        def __iter__(self):
+            for loader in self.loaders:
+                yield from loader
+
+
+    class CustomDataloaderTrainer(Trainer):
+        def get_train_dataloader(self):
+            dataloaders = [super(CustomDataloaderTrainer, self).get_train_dataloader() for _ in range(2)]
+            return MultiLoader(dataloaders)
+
+        def get_eval_dataloader(self, eval_dataset):
+            dataloaders = [super(CustomDataloaderTrainer, self).get_eval_dataloader(eval_dataset) for _ in range(2)]
+            return MultiLoader(dataloaders)
+
     class RegressionModel(nn.Module):
         def __init__(self, a=0, b=0, double_output=False):
             super().__init__()

From 5ee2d7b4330996382fda087d64b28010e5e76ba9 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 19:21:25 +0200
Subject: [PATCH 14/15] merge

---
 tests/trainer/test_trainer.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 040a7252a5de..99e18e7734f4 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -189,7 +189,6 @@ def __iter__(self):
                 yield self.dataset[self.current_sample]
                 self.current_sample += 1
 
-
     class MultiLoader:
         def __init__(self, loaders):
             self.loaders = loaders
@@ -201,7 +200,6 @@ def __iter__(self):
             for loader in self.loaders:
                 yield from loader
 
-
     class CustomDataloaderTrainer(Trainer):
         def get_train_dataloader(self):
             dataloaders = [super(CustomDataloaderTrainer, self).get_train_dataloader() for _ in range(2)]

From 1f8081f0f5eb9dce5310e59e3dea6aed81e026d8 Mon Sep 17 00:00:00 2001
From: Sander Land <sander@chatdesk.com>
Date: Tue, 29 Mar 2022 20:00:20 +0200
Subject: [PATCH 15/15] stylistic suggestion

---
 tests/trainer/test_trainer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 99e18e7734f4..afe97701d29b 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -202,11 +202,11 @@ def __iter__(self):
 
     class CustomDataloaderTrainer(Trainer):
         def get_train_dataloader(self):
-            dataloaders = [super(CustomDataloaderTrainer, self).get_train_dataloader() for _ in range(2)]
+            dataloaders = [super().get_train_dataloader(), super().get_train_dataloader()]
             return MultiLoader(dataloaders)
 
         def get_eval_dataloader(self, eval_dataset):
-            dataloaders = [super(CustomDataloaderTrainer, self).get_eval_dataloader(eval_dataset) for _ in range(2)]
+            dataloaders = [super().get_eval_dataloader(eval_dataset), super().get_eval_dataloader(eval_dataset)]
             return MultiLoader(dataloaders)
 
     class RegressionModel(nn.Module):