From e073d8497ddb440194eefc163eda31f525bfe06d Mon Sep 17 00:00:00 2001
From: Yash Raj Pandey <yashpn62@gmail.com>
Date: Wed, 15 Jul 2026 20:08:56 -0400
Subject: [PATCH] Cast float64 targets before moving them to the device

_predict_step_pytorch moved y to the model device first and only then cast
float64 down to float32. MPS does not support float64, so the move itself
raised before the cast could run:

  TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS
  framework doesn't support float64. Please use float32 instead.

numpy's default float dtype is float64, so an ordinary float target array
hits this on any Apple Silicon host. X_t already fuses the cast into its
own .to(device, dtype=torch.float32), so only y was affected.

Do the existing conditional cast on the CPU tensor and move afterwards.
The cast stays conditional on purpose: classifier targets are encoded as
int64 and regressor scaling can yield float16, so casting unconditionally
would silently change those dtypes.

Adds a regression test that drives the real _batch_forward path with
float64 targets on a model whose parameters live on MPS. It fails with the
TypeError above against the previous ordering. The test skips when MPS is
unavailable, since the crash only reproduces on that backend.

Signed-off-by: Yash Raj Pandey <yashpn62@gmail.com>
---
 tabfm/src/classifier_and_regressor.py         |  3 ++-
 .../classifier_and_regressor_pytorch_test.py  | 23 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index 4a0ebee..19fa61c 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -1977,9 +1977,10 @@ def _predict_step_pytorch(
   device = next(model.parameters()).device
 
   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)
 
   batch_size = X_batch.shape[0]
   train_size_t = torch.full(
diff --git a/tabfm/src/classifier_and_regressor_pytorch_test.py b/tabfm/src/classifier_and_regressor_pytorch_test.py
index 4ea1df0..62cac23 100644
--- a/tabfm/src/classifier_and_regressor_pytorch_test.py
+++ b/tabfm/src/classifier_and_regressor_pytorch_test.py
@@ -136,6 +136,29 @@ 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_batch_forward_float64_targets_on_mps(self):
+    if not torch.backends.mps.is_available():
+      self.skipTest("MPS is required for this test.")
+
+    class _Model(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.anchor = torch.nn.Parameter(torch.zeros((), device="mps"))
+
+      def forward(self, X, y, train_size, cat_mask=None, d=None):
+        del X, train_size, cat_mask, d
+        return y.unsqueeze(-1) + self.anchor
+
+    reg = TabFMRegressor(model=_Model(), batch_size=1)
+    X = np.zeros((1, 2, 1), dtype=np.float32)
+    y = np.zeros((1, 1), dtype=np.float64)
+
+    preds = reg._batch_forward(X, y)
+
+    self.assertEqual(preds.dtype, np.float32)
+    np.testing.assert_array_equal(preds, [[[-100.0]]])
+
 
 class PyTorchModelPickleTest(unittest.TestCase):
   """The PyTorch model must be picklable.