add batch

2024-05-05 19:28:40 +02:00 · 2024-05-05 19:28:40 +02:00 · c6f6ae28ca
parent 1cfb74db6a
commit c6f6ae28ca
1 changed files with 1 additions and 0 deletions
--- a/dataset/dataset.py
+++ b/dataset/dataset.py
@ -34,6 +34,7 @@ class Dataset:
        self.dataset = self.__load_dataset()\
            .shuffle(self.shuffle_buffer_size, seed=self.seed)\
            .repeat(self.repeat)\
+            .batch(self.batch_size, drop_remainder=True)\
            .prefetch(tf.data.experimental.AUTOTUNE)

    def __load_dataset(self) -> tf.data.Dataset: