diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index a785057..f0d2aa1 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -1819,7 +1819,7 @@ def _predict_step_pytorch(
   with torch.no_grad():
     out_t = model(X_t, y_t, train_size_t, cat_mask=cat_mask_t, d=d_t)
 
-  return out_t.cpu().numpy()
+  return out_t.float().cpu().numpy()  # upcast: numpy has no bfloat16
 
 
 # ---------------------------------------------------------------------------
diff --git a/tabfm/src/pytorch/model.py b/tabfm/src/pytorch/model.py
index 0accc02..8c47f7f 100644
--- a/tabfm/src/pytorch/model.py
+++ b/tabfm/src/pytorch/model.py
@@ -442,6 +442,11 @@ def __init__(self, *, embed_dim=8, max_classes=3, col_num_blocks=2,
                                     decoder_hidden or icl_dim * 2, is_classifier)
 
   def forward(self, x, y, train_size, cat_mask=None, d=None):
+    # Mirror the JAX model's entry: replace NaN with the -100 sentinel and cast
+    # to the compute dtype (JAX: `jnp.nan_to_num(X, nan=-100.0).astype(self.dtype)`).
+    # NaN is already imputed in the shared preprocessing, so nan_to_num is a
+    # no-op in the normal flow, but it keeps the model robust + JAX-faithful.
+    x = torch.nan_to_num(x, nan=-100.0).to(self.cls_tokens.dtype)
     emb = self.cell_embedder(x, y, train_size, cat_mask, d=d)
     emb = self.col_embedder(emb, train_size)
     b, t, _, e = emb.shape
diff --git a/tabfm/src/pytorch/tabfm_v1_0_0.py b/tabfm/src/pytorch/tabfm_v1_0_0.py
index 18c8258..88298d1 100644
--- a/tabfm/src/pytorch/tabfm_v1_0_0.py
+++ b/tabfm/src/pytorch/tabfm_v1_0_0.py
@@ -74,10 +74,20 @@ def load(
     checkpoint_path: Optional[str] = None,
     *,
     device: Optional[str] = None,
+    dtype: Any = torch.bfloat16,
     use_cache: bool = True,
 ) -> TabFM:
-  """Loads the PyTorch TabFM v1.0.0 model with pre-trained weights."""
-  cache_key = (model_type, checkpoint_path, device)
+  """Loads the PyTorch TabFM v1.0.0 model with pre-trained weights.
+
+  The checkpoint is stored in float32, but the model is designed to run in
+  bfloat16 (matching the JAX release's ``dtype=jnp.bfloat16`` compute default),
+  with a few internal fp32 upcasts. ``dtype`` casts the model accordingly; pass
+  ``None`` to keep the float32 weights.
+
+  ``dtype`` is provided for float32 debugging / quality comparison; the model is
+  designed for bfloat16 and this option may be removed in a future release.
+  """
+  cache_key = (model_type, checkpoint_path, device, str(dtype))
   
   if use_cache:
     _LOAD_CACHE_LOCK.acquire()
@@ -119,6 +129,9 @@ def load(
     state_dict = torch.load(checkpoint_file, map_location="cpu")
     model.load_state_dict(state_dict, strict=True)
 
+    if dtype is not None:
+      model = model.to(dtype)  # engage the bf16 compute design (see docstring)
+
     if device is not None:
       model = model.to(device)