Categorical bug fix

Author: snehilchatterjee

src/pytorch_tabular/categorical_encoders.py

@@ -62,6 +62,8 @@ def transform(self, X):
                 not X[self.cols].isnull().any().any()
             ), "`handle_missing` = `error` and missing values found in columns to encode."
         X_encoded = X.copy(deep=True)
+        category_cols = X_encoded.select_dtypes(include='category').columns
+        X_encoded[category_cols] = X_encoded[category_cols].astype('object')          
         for col, mapping in self._mapping.items():
             X_encoded[col] = X_encoded[col].fillna(NAN_CATEGORY).map(mapping["value"])
 
@@ -267,4 +269,4 @@ def save_as_object_file(self, path):
 
     def load_from_object_file(self, path):
         for k, v in pickle.load(open(path, "rb")).items():
-            setattr(self, k, v)
+            setattr(self, k, v)

src/pytorch_tabular/tabular_datamodule.py

@@ -301,10 +301,14 @@ def _update_config(self, config) -> InferredConfig:
         else:
             raise ValueError(f"{config.task} is an unsupported task.")
         if self.train is not None:
+            category_cols = self.train[config.categorical_cols].select_dtypes(include='category').columns
+            self.train[category_cols] = self.train[category_cols].astype('object')
             categorical_cardinality = [
                 int(x) + 1 for x in list(self.train[config.categorical_cols].fillna("NA").nunique().values)
             ]
         else:
+            category_cols = self.train_dataset.data[config.categorical_cols].select_dtypes(include='category').columns
+            self.train_dataset.data[category_cols] = self.train_dataset.data[category_cols].astype('object')
             categorical_cardinality = [
                 int(x) + 1 for x in list(self.train_dataset.data[config.categorical_cols].nunique().values)
             ]
@@ -805,6 +809,7 @@ def train_dataloader(self, batch_size: Optional[int] = None) -> DataLoader:
             num_workers=self.config.num_workers,
             sampler=self.train_sampler,
             pin_memory=self.config.pin_memory,
+            **self.config.dataloader_kwargs,
         )
 
     def val_dataloader(self, batch_size: Optional[int] = None) -> DataLoader:
@@ -823,6 +828,7 @@ def val_dataloader(self, batch_size: Optional[int] = None) -> DataLoader:
             shuffle=False,
             num_workers=self.config.num_workers,
             pin_memory=self.config.pin_memory,
+            **self.config.dataloader_kwargs,
         )
 
     def _prepare_inference_data(self, df: DataFrame) -> DataFrame:
@@ -865,6 +871,7 @@ def prepare_inference_dataloader(
             batch_size or self.batch_size,
             shuffle=False,
             num_workers=self.config.num_workers,
+            **self.config.dataloader_kwargs,
         )
 
     def save_dataloader(self, path: Union[str, Path]) -> None: