fix device bug on GPU (#254)

Author: noahho

src/tabpfn/base.py

@@ -28,7 +28,7 @@
     InferenceEngineOnDemand,
 )
 from tabpfn.model.loading import load_model_criterion_config
-from tabpfn.utils import infer_fp16_inference_mode
+from tabpfn.utils import infer_device_and_type, infer_fp16_inference_mode
 
 if TYPE_CHECKING:
     import numpy as np
@@ -254,14 +254,15 @@ def check_cpu_warning(
         X: The input data (NumPy array, Pandas DataFrame, or Torch Tensor)
     """
     allow_cpu_override = os.getenv("TABPFN_ALLOW_CPU_LARGE_DATASET", "0") == "1"
+    device_mapped = infer_device_and_type(device)
 
     # Determine number of samples
     try:
         num_samples = X.shape[0]
     except AttributeError:
         return
 
-    if device == torch.device("cpu") or device == "cpu" or "cpu" in device:
+    if torch.device(device_mapped).type == "cpu":
         if num_samples > 1000:
             if not allow_cpu_override:
                 raise RuntimeError(