Update train.py

train.py

@@ -36,14 +36,10 @@ def caption_batch(batch, processor, model):
 
     pil_images = []
     for image in images:
-        if isinstance(image, torch.Tensor):
-            image = image.cpu().numpy()
-        
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        pil_images.append(image)
+        if isinstance(image, Image.Image):
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            pil_images.append(image)
 
     msg = [
         {
@@ -110,9 +106,6 @@ def process_shard_worker(
     else:
         shard = cast(Dataset, loaded)
 
-    shard = shard.with_format("torch")
-    shard.set_format(type="torch", columns=["image"])
-
     print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
     result = shard.map(
         lambda batch: caption_batch(batch, processor, model),
@@ -120,7 +113,6 @@ def process_shard_worker(
         batch_size=batch_size,
         remove_columns=[col for col in shard.column_names if col != "image"],
         writer_batch_size=1000,
-        keep_in_memory=True,
     )
 
     print(f"[GPU {gpu_id}] Saving results to {output_file}...", flush=True)