remove unused generator argument

8ed52b70 · Kolja Stahl · ceec89a6 · 8ed52b70
Commit 8ed52b70 authored Jan 20, 2022 by Kolja Stahl
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

openfold/data/data_modules.py openfold/data/data_modules.py +2 -2

No files found.
--- a/openfold/data/data_modules.py
+++ b/openfold/data/data_modules.py
@@ -245,7 +245,7 @@ class OpenFoldDataset(torch.utils.data.IterableDataset):


 class OpenFoldBatchCollator:
-    def __init__(self, config, generator, stage="train"):
+    def __init__(self, config, stage="train"):
        self.stage = stage
        self.feature_pipeline = feature_pipeline.FeaturePipeline(config)

@@ -503,7 +503,7 @@ class OpenFoldDataModule(pl.LightningDataModule):
        else:
            raise ValueError("Invalid stage")

-        batch_collator = OpenFoldBatchCollator(self.config, generator, stage)
+        batch_collator = OpenFoldBatchCollator(self.config, stage)

        dl = OpenFoldDataLoader(
            dataset,