updated again idk what

train_llama.py

@@ -59,7 +59,9 @@ def tokenize_data(example):
     }
 
 tokenized_dataset = dataset["train"].map(tokenize_data, batched=False, remove_columns=dataset["train"].column_names)
-print("First tokenized example:", {k: (type(v), v.shape) for k, v in tokenized_dataset[0].items()})
+# Fix print to handle potential list or tensor
+first_example = tokenized_dataset[0]
+print("First tokenized example:", {k: (type(v), v.shape if hasattr(v, 'shape') else len(v)) for k, v in first_example.items()})
 
 # Data collator
 def custom_data_collator(features):