diff --git a/src/python/setup.py b/src/python/setup.py
index f29b3c19..35b3d89c 100644
--- a/src/python/setup.py
+++ b/src/python/setup.py
@@ -66,3 +66,4 @@
     package_data={"tensorflow_cloud": ["tuner/api/*.json"]},
     packages=find_packages(where=relative_directory),
 )
+train_dataset = train_dataset.batch(BATCH_SIZE).map(scale)