Add **dataloader_args to dataloaders in torch trainer

epochalyst/pipeline/model/training/torch_trainer.py

@@ -142,6 +142,8 @@ def log_to_terminal(self, message: str) -> None:
     _fold: int = field(default=-1, init=False, repr=False, compare=False)
     n_folds: float = field(default=-1, init=True, repr=False, compare=False)
 
+    dataloader_args: dict[str, Any] = field(default_factory=dict, repr=False)
+
     def __post_init__(self) -> None:
         """Post init method for the TorchTrainer class."""
         # Make sure to_predict is either "test" or "all" or "none"
@@ -397,6 +399,7 @@ def predict_on_loader(
             collate_fn=(
                 collate_fn if hasattr(loader.dataset, "__getitems__") else None  # type: ignore[arg-type]
             ),
+            **self.dataloader_args,
         )
         with torch.no_grad(), tqdm(loader, unit="batch", disable=False) as tepoch:
             for data in tepoch:
@@ -473,12 +476,14 @@ def create_dataloaders(
             batch_size=self.batch_size,
             shuffle=True,
             collate_fn=(collate_fn if hasattr(train_dataset, "__getitems__") else None),  # type: ignore[arg-type]
+            **self.dataloader_args,
         )
         test_loader = DataLoader(
             test_dataset,
             batch_size=self.batch_size,
             shuffle=False,
             collate_fn=(collate_fn if hasattr(test_dataset, "__getitems__") else None),  # type: ignore[arg-type]
+            **self.dataloader_args,
         )
         return train_loader, test_loader
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "epochalyst"
-version = "0.3.2"
+version = "0.3.3"
 authors = [
     { name = "Jasper van Selm", email = "jmvanselm@gmail.com" },
     { name = "Ariel Ebersberger", email = "arielebersberger@gmail.com" },