Update data.py

Author: zsdonghao

data.py

@@ -47,7 +47,7 @@ def _map_fn(image_path):
     train_ds = tf.data.Dataset.from_generator(generator_train, output_types=tf.string)
     ds = train_ds.shuffle(buffer_size=4096)
     # ds = ds.shard(num_shards=hvd.size(), index=hvd.rank())
-    ds = ds.repeat(n_epoch)
+    # ds = ds.repeat(n_epoch)
     ds = ds.map(_map_fn, num_parallel_calls=4)
     ds = ds.batch(batch_size)
     ds = ds.prefetch(buffer_size=2)