diff --git a/dataset_loader.py b/dataset_loader.py
index 5cc74e1..e3f0519 100644
--- a/dataset_loader.py
+++ b/dataset_loader.py
@@ -93,10 +93,10 @@ def load_single_image(img_path):
 	return input_pair(img_id, src_image, trt_image, rotation, translation, fov, rotation_pred)
 
   ds = tf.data.Dataset.list_files(os.path.join(data_path, '*'))
-  ds = ds.flat_map(load_sstable)
+  ds = ds.flat_map(load_data)
   ds = ds.map(load_images, num_parallel_calls=50).apply(
   	tf.data.experimental.ignore_errors()).repeat(epochs)
   ds = ds.batch(batch_size, drop_remainder=True).prefetch(10)
   return ds
 
-  
\ No newline at end of file
+