From 178a4523ef06e272128e431d9d8adad1fa1c7592 Mon Sep 17 00:00:00 2001 From: shahules786 Date: Fri, 21 Oct 2022 09:48:28 +0530 Subject: [PATCH] fix worker init fn --- enhancer/data/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/enhancer/data/dataset.py b/enhancer/data/dataset.py index 06a6f67..218cd7c 100644 --- a/enhancer/data/dataset.py +++ b/enhancer/data/dataset.py @@ -189,8 +189,8 @@ class TaskDataset(pl.LightningDataModule): worker_info = torch.utils.data.get_worker_info() dataset = worker_info.dataset worker_id = worker_info.id - split_size = len(dataset.data) // worker_info.num_workers - dataset.data = dataset.data[ + split_size = len(dataset.dataset.train_data) // worker_info.num_workers + dataset.data = dataset.dataset.train_data[ worker_id * split_size : (worker_id + 1) * split_size ]