From fcc09e70d61fb6f0164a03ece89f8aca044e615a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Koz=C5=82owski?=
 <michalkozlowski936@gmail.com>
Date: Tue, 24 Jan 2023 18:21:55 +0100
Subject: [PATCH] t

---
 utils/donut_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/donut_dataset.py b/utils/donut_dataset.py
index 4db3a1e..3c4651c 100644
--- a/utils/donut_dataset.py
+++ b/utils/donut_dataset.py
@@ -47,7 +47,7 @@ class DonutDataset(Dataset):
         self.sort_json_key = sort_json_key
         self.added_tokens = added_tokens
 
-        self.dataset = dataset
+        self.dataset = dataset.with_format("torch")
         # self.dataset_length = len(self.dataset.with_format("torch"))
 
         self.gt_token_sequences = []