Add pickle_protocol to DataConfig for passing

it down to `torch.save()` when caching datasets to disk

Author: sorenmacbeth

src/pytorch_tabular/config/config.py

@@ -96,6 +96,8 @@ class DataConfig:
         handle_missing_values (bool): Whether to handle missing values in categorical columns as
                 unknown
 
+        pickle_protocol (int): pickle protocol version passed to `torch.save` for dataset caching to disk
+
         dataloader_kwargs (Dict[str, Any]): Additional kwargs to be passed to PyTorch DataLoader. See
                 https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
 
@@ -179,6 +181,11 @@ class DataConfig:
         metadata={"help": "Whether or not to handle missing values in categorical columns as unknown"},
     )
 
+    pickle_protocol: int = field(
+        default=2,
+        metadata={"help": "pickle protocol version passed to `torch.save` for dataset caching to disk"},
+    )
+
     dataloader_kwargs: Dict[str, Any] = field(
         default_factory=dict,
         metadata={"help": "Additional kwargs to be passed to PyTorch DataLoader."},

src/pytorch_tabular/tabular_datamodule.py

@@ -484,8 +484,10 @@ def _cache_dataset(self):
         self.validation = None
 
         if self.cache_mode is self.CACHE_MODES.DISK:
-            torch.save(train_dataset, self.cache_dir / "train_dataset")
-            torch.save(validation_dataset, self.cache_dir / "validation_dataset")
+            torch.save(train_dataset, self.cache_dir / "train_dataset", pickle_protocol=self.config.pickle_protocol)
+            torch.save(
+                validation_dataset, self.cache_dir / "validation_dataset", pickle_protocol=self.config.pickle_protocol
+            )
         elif self.cache_mode is self.CACHE_MODES.MEMORY:
             self.train_dataset = train_dataset
             self.validation_dataset = validation_dataset