Update dataset.py

Author: MaxSchambach

mambular/data_utils/dataset.py

@@ -44,7 +44,8 @@ def __init__(
             self.labels = None  # No labels in prediction mode
 
     def __len__(self):
-        return len(self.num_features_list[0])  # Use numerical features length
+        _feats = self.num_features_list if self.num_features_list else cat_features_list
+        return len(_feats[0])
 
     def __getitem__(self, idx):
         """Retrieves the features and label for a given index.