Update generation.py

generation.py

@@ -1197,8 +1197,9 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             self.vad_seek_callback(kwargs["stno_mask"])
         if "is_valid" in kwargs:
             kwargs['is_valid'] = kwargs["is_valid"][batch_idx_map]
-        kwargs['labels'] = kwargs["labels"][batch_idx_map]
-        kwargs['upp_labels'] = kwargs["upp_labels"][batch_idx_map]
+        if "labels" in kwargs:
+            kwargs['labels'] = kwargs["labels"][batch_idx_map]
+            kwargs['upp_labels'] = kwargs["upp_labels"][batch_idx_map]
         return kwargs
 
     def generate_with_fallback(