Update dataset.py

dataset.py

@@ -70,7 +70,7 @@ class DemoDataset(Dataset):
             )
             prompt_features.append(prompt_feature)
         
-        prompt_labels = torch.tensor(self.demonstration_labels, dtype=torch.long)
+        prompt_labels = torch.tensor([self.demonstration_labels], dtype=torch.long)
         
         return {
             'main_features': main_features,
@@ -113,7 +113,8 @@ def collate_fn(batch):
     file_names = [item['file_name'] for item in batch]
     file_paths = [item['file_path'] for item in batch]
 
-    prompt_labels = torch.tensor([item['prompt_labels'] for item in batch], dtype=torch.long)
+    # 确保 prompt_labels 的形状正确 [batch_size, num_prompts]
+    prompt_labels = torch.cat([item['prompt_labels'] for item in batch], dim=0)
 
     return {
         'main_features': main_features,