Reinstantiate weighted captions after a necessary revert to Main

library/train_util.py

@@ -950,10 +950,10 @@ class BaseDataset(torch.utils.data.Dataset):
         example["images"] = images
 
         example["latents"] = torch.stack(latents_list) if latents_list[0] is not None else None
+        example["captions"] = captions
 
         if self.debug_dataset:
             example["image_keys"] = bucket[image_index : image_index + self.batch_size]
-            example["captions"] = captions
         return example
 
 
@@ -3097,4 +3097,4 @@ class collater_class:
         # set epoch and step
         dataset.set_current_epoch(self.current_epoch.value)
         dataset.set_current_step(self.current_step.value)
-        return examples[0]
+        return examples[0]