Update train.py

train.py

@@ -122,10 +122,10 @@ def process_shard_worker(
 def main():
     import multiprocessing as mp
 
-    input_dataset = "nroggendorff/fries"
-    output_dataset = "nroggendorff/fries"
+    input_dataset = "none-yet/anime-captions"
+    output_dataset = "none-yet/anime-captions"
     model_name = "datalab-to/chandra"
-    batch_size = 2
+    batch_size = 32
 
     print("Loading dataset info...")
     loaded = datasets.load_dataset(input_dataset, split="train")