Update train.py

train.py

@@ -36,6 +36,9 @@ 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":