fix_base_model

Nixtla · Nov 19, 2024 · b2c7691 · b2c7691
1 parent c529ced
commit b2c7691
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/nbs/common.base_model.ipynb b/nbs/common.base_model.ipynb
@@ -158,6 +158,7 @@
     "        optimizer_kwargs: Union[Dict, None] = None,\n",
     "        lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, None] = None,\n",
     "        lr_scheduler_kwargs: Union[Dict, None] = None,\n",
+    "        dataloader_kwargs=None,\n",
     "        **trainer_kwargs,\n",
     "    ):\n",
     "        super().__init__()\n",
@@ -364,6 +365,7 @@
     "\n",
     "        # DataModule arguments\n",
     "        self.num_workers_loader = num_workers_loader\n",
+    "        self.dataloader_kwargs = dataloader_kwargs\n",
     "        self.drop_last_loader = drop_last_loader\n",
     "        # used by on_validation_epoch_end hook\n",
     "        self.validation_step_outputs: List = []\n",

diff --git a/neuralforecast/common/_base_model.py b/neuralforecast/common/_base_model.py
@@ -111,6 +111,7 @@ def __init__(
         optimizer_kwargs: Union[Dict, None] = None,
         lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, None] = None,
         lr_scheduler_kwargs: Union[Dict, None] = None,
+        dataloader_kwargs=None,
         **trainer_kwargs,
     ):
         super().__init__()
@@ -352,6 +353,7 @@ def __init__(
 
         # DataModule arguments
         self.num_workers_loader = num_workers_loader
+        self.dataloader_kwargs = dataloader_kwargs
         self.drop_last_loader = drop_last_loader
         # used by on_validation_epoch_end hook
         self.validation_step_outputs: List = []