diff --git a/enhancer/data/dataset.py b/enhancer/data/dataset.py
index 8110a2a..06a6f67 100644
--- a/enhancer/data/dataset.py
+++ b/enhancer/data/dataset.py
@@ -185,7 +185,7 @@ class TaskDataset(pl.LightningDataModule):
         output["noisy"] = torch.stack(output["noisy"], dim=0)
         return output
 
-    def worker_init_fn(self):
+    def worker_init_fn(self, _):
         worker_info = torch.utils.data.get_worker_info()
         dataset = worker_info.dataset
         worker_id = worker_info.id