From 821f60541ea42b5f9f437e5f41bc0985305a463d Mon Sep 17 00:00:00 2001 From: jhauret Date: Tue, 18 Jun 2024 16:48:15 +0200 Subject: [PATCH] try to pass setup when "train" and "validation" are not defined --- vibravox/lightning_datamodules/spkv.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/vibravox/lightning_datamodules/spkv.py b/vibravox/lightning_datamodules/spkv.py index 467fba8..84a4a18 100644 --- a/vibravox/lightning_datamodules/spkv.py +++ b/vibravox/lightning_datamodules/spkv.py @@ -84,7 +84,7 @@ def setup(self, stage=None): if train_dataset is None or val_dataset is None: pass - if self.sensorA == self.sensorB: + if self.sensorA == self.sensorB and train_dataset is not None and val_dataset is not None: # When self.sensorA and self.sensorB are the same, only generate the dataset using one column # Only keep the relevant columns for this task : @@ -109,7 +109,7 @@ def setup(self, stage=None): train_dataset = train_dataset.rename_column(f"audio.{self.sensorA}", "audio") val_dataset = val_dataset.rename_column(f"audio.{self.sensorA}", "audio") - else: + elif self.sensorA != self.sensorB and train_dataset is not None and val_dataset is not None: # When self.sensorA and self.sensorB are different, generate the dataset by interleaving both sensors # in the same dataset for train/validation, which results in a two times larger dataset for training, # but allows to learn embeddings for both sensors @@ -164,9 +164,11 @@ def setup(self, stage=None): val_dataset = interleave_datasets(datasets=[val_dataset_a, val_dataset_b], probabilities=[0.5, 0.5], stopping_strategy='all_exhausted') - # Setting format to torch - self.train_dataset = train_dataset.with_format("torch") - self.val_dataset = val_dataset.with_format("torch") + + if train_dataset is not None and val_dataset is not None: + # Setting format to torch + self.train_dataset = train_dataset.with_format("torch") + self.val_dataset = val_dataset.with_format("torch") if stage == "test": # Generating dataset for testing for Speaker Verification (only for the test set) : pairs are needed