diff --git a/worker_generic_training_dataset/worker.py b/worker_generic_training_dataset/worker.py index 7270050fb47d8cd43c5cf218789d9dfc21d74104..30120a142d8c518169d5c14acecafb46f266446f 100644 --- a/worker_generic_training_dataset/worker.py +++ b/worker_generic_training_dataset/worker.py @@ -76,15 +76,15 @@ class DatasetExtractor(BaseWorker): """ logger.info("Retrieving information from process_information") - train_folder_id = self.process_information.get("train_folder_id") + train_folder_id = self.config.get("train_folder_id") assert train_folder_id, "A training folder id is necessary to use this worker" self.training_folder_id = UUID(train_folder_id) - val_folder_id = self.process_information.get("validation_folder_id") + val_folder_id = self.config.get("validation_folder_id") assert val_folder_id, "A validation folder id is necessary to use this worker" self.validation_folder_id = UUID(val_folder_id) - test_folder_id = self.process_information.get("test_folder_id") + test_folder_id = self.config.get("test_folder_id") self.testing_folder_id = UUID(test_folder_id) if test_folder_id else None def initialize_database(self):