From 587197dcd2b50ad9e96aedbfa389bf4fcc294c3c Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Mon, 15 Feb 2021 05:49:07 -0500
Subject: [PATCH] Fix datasets set_format (#10178)

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

diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 858622d09e..b0f7e694ea 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -439,7 +439,8 @@ class Trainer:
             f"The following columns {dset_description}don't have a corresponding argument in "
             f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
         )
-        dataset.set_format(type=dataset.format["type"], columns=columns)
+
+        dataset.set_format(type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"])
 
     def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
         if isinstance(self.train_dataset, torch.utils.data.IterableDataset) or not isinstance(