Update tabular_datamodule.py for torch.load

ArozHada · web-flow · commit e59e21e8da0e · 2025-02-16T15:37:09.000+01:00
diff --git a/src/pytorch_tabular/tabular_datamodule.py b/src/pytorch_tabular/tabular_datamodule.py
@@ -758,7 +758,7 @@ def _load_dataset_from_cache(self, tag: str = "train"):
                 )
         elif self.cache_mode is self.CACHE_MODES.DISK:
             try:
-                dataset = torch.load(self.cache_dir / f"{tag}_dataset")
+                dataset = torch.load(self.cache_dir / f"{tag}_dataset", weights_only=False)
             except FileNotFoundError:
                 raise FileNotFoundError(
                     f"{tag}_dataset not found in {self.cache_dir}. Please provide the" f" data for {tag} dataloader"

Original file line number	Diff line number	Diff line change
`@@ -758,7 +758,7 @@ def _load_dataset_from_cache(self, tag: str = "train"):`
`758`	`758`	`)`
`759`	`759`	`elif self.cache_mode is self.CACHE_MODES.DISK:`
`760`	`760`	`try:`
`761`		`- dataset = torch.load(self.cache_dir / f"{tag}_dataset")`
	`761`	`+ dataset = torch.load(self.cache_dir / f"{tag}_dataset", weights_only=False)`
`762`	`762`	`except FileNotFoundError:`
`763`	`763`	`raise FileNotFoundError(`
`764`	`764`	`f"{tag}_dataset not found in {self.cache_dir}. Please provide the" f" data for {tag} dataloader"`