fix index error in cls token

basf · May 31, 2024 · 4acaf92 · 4acaf92
1 parent 141c595
commit 4acaf92
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/mambular/base_models/classifier.py b/mambular/base_models/classifier.py
@@ -232,9 +232,7 @@ def forward(self, num_features, cat_features):
             The output predictions of the model.
         """
         batch_size = (
-            cat_features[0].size(0)
-            if cat_features is not None
-            else num_features[0].size(0)
+            cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
         )
         cls_tokens = self.cls_token.expand(batch_size, -1, -1)