Cast float64 targets before the device move in the context-cache path

tabfm/src/classifier_and_regressor.py

@@ -2064,9 +2064,10 @@ def _build_context_cache_pytorch(
         Xs_split, ys_split, cat_masks_split, ds_split
     ):
       X_t = torch.from_numpy(X_batch).to(device, dtype=torch.float32)
-      y_t = torch.from_numpy(y_batch).to(device)
+      y_t = torch.from_numpy(y_batch)
       if y_t.dtype == torch.float64:
         y_t = y_t.to(torch.float32)
+      y_t = y_t.to(device)
       cat_mask_t = torch.from_numpy(cat_mask_batch).to(device)
       d_t = torch.from_numpy(ds_batch).to(device)
       _, cache = model.prefill(X_t, y_t, cat_mask=cat_mask_t, d=d_t)

tabfm/src/classifier_and_regressor_pytorch_test.py

@@ -136,6 +136,49 @@ def test_regressor_fit_predict(self):
     self.assertEqual(preds_cached.shape, (10,))
     np.testing.assert_allclose(preds_cached, preds, rtol=1e-5, atol=1e-6)
 
+  def test_regressor_cache_context_float64_targets_on_mps(self):
+    # _build_context_cache_pytorch moved y to the device before casting
+    # float64 down to float32. MPS rejects float64 tensors at transfer time,
+    # so fit(cache_context=True) crashed on Apple Silicon for any float64
+    # target array (numpy's default float dtype). CPU/CUDA accept float64,
+    # so the ordering is only observable on MPS.
+    if not torch.backends.mps.is_available():
+      self.skipTest("MPS is required for this test.")
+
+    np.random.seed(42)
+    model = pytorch_model.TabFM(
+        embed_dim=8,
+        max_classes=1,
+        col_num_blocks=1,
+        col_nhead=2,
+        col_num_inds=8,
+        row_num_blocks=1,
+        row_nhead=2,
+        row_num_cls=2,
+        icl_num_blocks=1,
+        icl_nhead=2,
+        ff_factor=2,
+        feature_group_size=2,
+        is_classifier=False
+    ).to("mps")
+
+    reg = TabFMRegressor(
+        model=model,
+        n_estimators=2,
+        batch_size=2,
+        random_state=42,
+        cache_context=True,
+        maybe_quantize_kv_cache=False,
+    )
+
+    X = np.random.rand(10, 3)
+    y = np.random.rand(10)  # float64 — numpy's default float dtype
+
+    reg.fit(X, y)
+    preds = reg.predict(X)
+    self.assertEqual(preds.shape, (10,))
+    self.assertTrue(np.all(np.isfinite(preds)))
+
 
 class PyTorchModelPickleTest(unittest.TestCase):
   """The PyTorch model must be picklable.