(#4) Refactor: Fix typo

betarixm · Mar 30, 2022 · 6d137bf · 6d137bf
1 parent 4ae24a7
commit 6d137bf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/typings/dataset.py b/src/typings/dataset.py
@@ -56,7 +56,7 @@ def dataset(
             .batch(batch)
         )
 
-        test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
+        test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch)
 
         return train_ds, test_ds