Update Mo.py

Mo.py

@@ -99,6 +99,8 @@ def txt_stream(file_path):
             )
 
 
+LIMIT = 500000  # 원하는 만큼
+
 dataset = tf.data.Dataset.from_generator(
     lambda: txt_stream(DATA_PATH),
     output_signature=(
@@ -107,7 +109,7 @@ dataset = tf.data.Dataset.from_generator(
     )
 )
 
-dataset = dataset.shuffle(2000, seed=SEED).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
+dataset = dataset.take(LIMIT).shuffle(2000, seed=SEED).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
 
 with strategy.scope():
     dist_dataset = strategy.experimental_distribute_dataset(dataset)