fix: dataset update

piercus · Jan 26, 2024 · f96d8e9 · f96d8e9
1 parent 3bb845d
commit f96d8e9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/refiners/training_utils/datasets/color_palette.py b/src/refiners/training_utils/datasets/color_palette.py
@@ -83,7 +83,7 @@ def get_color_palette(self, index: int) -> ColorPalette:
         probabilities = weights / sum
         palette_index = int(random.choices(choices, probabilities, k=1)[0])
         item = self.hf_dataset[index]
-        palette: ColorPalette = item[f"palette_{palette_index}"]
+        palette: ColorPalette = item[f"palettes"][str(palette_index)]
         return palette
 
     def collate_fn(self, batch: list[TextEmbeddingColorPaletteLatentsBatch]) -> TextEmbeddingColorPaletteLatentsBatch: