Update dataloader.py

splendor1811 · Jun 12, 2023 · 3de95b9 · 3de95b9
1 parent 238d029
commit 3de95b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dataloader/dataloader.py b/dataloader/dataloader.py
@@ -122,11 +122,11 @@ def build_dataset(paths, labels=None, bsize=32, cache=True,
  dset = dset.map(decode_fn, num_parallel_calls=AUTO)
  dset = dset.cache(cache_dir) if cache else dset
  dset = dset.map(augment_fn, num_parallel_calls=AUTO) if augment else dset
- dset = dset.repeat() if repeat else dset
+ #dset = dset.repeat() if repeat else dset
  dset = dset.shuffle(shuffle) if shuffle else dset
  dset = dset.batch(bsize)
  # dset = dset.map(BatchAdvAugment, num_parallel_calls=AUTO) if augmentAdv else dset
  dset = dset.map(BatchAdvAugmentSeg, num_parallel_calls=AUTO) if augmentAdvSeg else dset
  dset = dset.prefetch(AUTO)
 
- return dset
+ return dset