restore num_workers in predict

Nixtla · Nov 8, 2024 · 0ab3839 · 0ab3839
1 parent 9e86685
commit 0ab3839
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/nbs/common.base_recurrent.ipynb b/nbs/common.base_recurrent.ipynb
@@ -553,6 +553,7 @@
     "        datamodule = TimeSeriesDataModule(\n",
     "            dataset=dataset,\n",
     "            valid_batch_size=self.valid_batch_size,\n",
+    "            num_workers=self.num_workers_loader,\n",
     "            **data_module_kwargs\n",
     "        )\n",
     "        fcsts = trainer.predict(self, datamodule=datamodule)\n",

diff --git a/neuralforecast/common/_base_recurrent.py b/neuralforecast/common/_base_recurrent.py
@@ -575,6 +575,7 @@ def predict(self, dataset, step_size=1, random_seed=None, **data_module_kwargs):
         datamodule = TimeSeriesDataModule(
             dataset=dataset,
             valid_batch_size=self.valid_batch_size,
+            num_workers=self.num_workers_loader,
             **data_module_kwargs,
         )
         fcsts = trainer.predict(self, datamodule=datamodule)