adjust order in __getitem__ functionality and batch for lightningmodule

Author: AnFreTh

mambular/base_models/lightning_wrapper.py

@@ -188,7 +188,7 @@ def training_step(self, batch, batch_idx):  # type: ignore
         Tensor
             Training loss.
         """
-        cat_features, num_features, labels = batch
+        num_features, cat_features, labels = batch
 
         # Check if the model has a `penalty_forward` method
         if hasattr(self.base_model, "penalty_forward"):
@@ -235,7 +235,7 @@ def validation_step(self, batch, batch_idx):  # type: ignore
             Validation loss.
         """
 
-        cat_features, num_features, labels = batch
+        num_features, cat_features, labels = batch
         preds = self(num_features=num_features, cat_features=cat_features)
         val_loss = self.compute_loss(preds, labels)
 
@@ -277,7 +277,7 @@ def test_step(self, batch, batch_idx):  # type: ignore
         Tensor
             Test loss.
         """
-        cat_features, num_features, labels = batch
+        num_features, cat_features, labels = batch
         preds = self(num_features=num_features, cat_features=cat_features)
         test_loss = self.compute_loss(preds, labels)
 
@@ -308,7 +308,7 @@ def predict_step(self, batch, batch_idx):
             Predictions.
         """
 
-        cat_features, num_features = batch
+        num_features, cat_features = batch
         preds = self(num_features=num_features, cat_features=cat_features)
 
         return preds

mambular/data_utils/dataset.py

@@ -20,7 +20,9 @@ class MambularDataset(Dataset):
         regression (bool, optional): A flag indicating if the dataset is for a regression task. Defaults to True.
     """
 
-    def __init__(self, cat_features_list, num_features_list, labels=None, regression=True):
+    def __init__(
+        self, cat_features_list, num_features_list, labels=None, regression=True
+    ):
         self.cat_features_list = cat_features_list  # Categorical features tensors
         self.num_features_list = num_features_list  # Numerical features tensors
         self.regression = regression
@@ -54,7 +56,9 @@ def __getitem__(self, idx):
             tuple: A tuple containing two lists of tensors (one for categorical features and one for numerical features)
             and a single label (if available).
         """
-        cat_features = [feature_tensor[idx] for feature_tensor in self.cat_features_list]
+        cat_features = [
+            feature_tensor[idx] for feature_tensor in self.cat_features_list
+        ]
         num_features = [
             torch.as_tensor(feature_tensor[idx]).clone().detach().to(torch.float32)
             for feature_tensor in self.num_features_list
@@ -68,7 +72,6 @@ def __getitem__(self, idx):
                 label = label.clone().detach().to(torch.float32)
             else:
                 label = label.clone().detach().to(torch.long)
-            return cat_features, num_features, label
+            return num_features, cat_features, label
         else:
-            return cat_features, num_features  # No label in prediction mode
-
+            return num_features, cat_features  # No label in prediction mode