Cast float64 targets before moving them to the device

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(

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.