Update Training script

train2.py

@@ -38,11 +38,12 @@ print(f"🔹 Loading processor from {MODEL_NAME}...")
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 
 def transform(example):
-    # Determine image column
-    image_column = "image" if "image" in example else list(example.keys())[0]
+    # Determine correct image column
+    image_column = "image" if "image" in example else [c for c in example.keys() if c != "label"][0]
+
     images = example[image_column]
 
-    # Ensure list
+    # Ensure we always have a list
     if not isinstance(images, list):
         images = [images]
 
@@ -56,7 +57,7 @@ def transform(example):
             raise ValueError(f"Unknown type for image: {type(img)}")
         processed_images.append(img)
 
-    # Convert images to tensors
+    # Convert to tensors (batched)
     inputs = processor(images=processed_images, return_tensors="pt")
 
     # Handle labels