fixed device bug in TabM for GPU

Author: dholzmueller

README.md

@@ -196,6 +196,7 @@ and https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html
 
 ## Releases (see git tags)
 
+- v1.5.1: fixed a device bug in TabM for GPU
 - v1.5.0:
     - added `n_repeats` parameter to scikit-learn interfaces for repeated cross-validation
     - HPO sklearn interfaces (the ones using random search)

pytabkit/__about__.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "1.5.0"
+__version__ = "1.5.1"

pytabkit/models/alg_interfaces/tabm_interface.py

@@ -261,7 +261,7 @@ def evaluate(part: str) -> float:
                         eval_batch_size
                     )
                     ]
-                )
+                ).cpu()
             )
             if task_type == 'regression':
                 # Transform the predictions back to the original label space.
@@ -277,7 +277,7 @@ def evaluate(part: str) -> float:
             if not average_logits:
                 y_pred = y_pred.mean(dim=1)
 
-            y_true = data[part]['y']
+            y_true = data[part]['y'].cpu()
             if task_type == 'regression' and len(y_true.shape) == 1:
                 y_true = y_true.unsqueeze(-1)
             if task_type == 'regression' and len(y_pred.shape) == 1: