Fix data loader init

borchero · Dec 21, 2021 · f12731e · f12731e
1 parent 9ddc139
commit f12731e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lightkit/data/loader.py b/lightkit/data/loader.py
@@ -42,6 +42,9 @@ def __init__(self, dataset: Dataset[T_co], **kwargs: Any):
                 batch_size=kwargs.get("batch_size", 1),
                 drop_last=kwargs.get("drop_last", False),
             )
+            kwargs.pop("batch_size", None)
+            kwargs.pop("shuffle", None)
+            kwargs.pop("drop_last", None)
             kwargs.setdefault("collate_fn", collate_tuple)
 
         super().__init__(dataset, **kwargs)  # type: ignore