Stabilize classifier logits for inference

scripts/tmos_classifier.py

@@ -72,8 +72,8 @@ class TMOSClassifier(nn.Module):
         for param in self.base.lm_head.parameters():
             param.requires_grad = False
 
-        # Binary classification head: hidden_size → 1 logit
-        self.classifier = nn.Linear(hidden_size, 1, dtype=torch_dtype)
+        # Keep the classifier head in fp32 for numerical stability.
+        self.classifier = nn.Linear(hidden_size, 1, dtype=torch.float32)
         nn.init.xavier_uniform_(self.classifier.weight)
         nn.init.zeros_(self.classifier.bias)
 
@@ -122,8 +122,16 @@ class TMOSClassifier(nn.Module):
             # No mask → just take the last position
             pooled = last_hidden_state[:, -1, :]
 
+        # Replace non-finite activations defensively before the classifier.
+        pooled = torch.nan_to_num(pooled, nan=0.0, posinf=1e4, neginf=-1e4)
+
+        # Match classifier device to pooled features when model is sharded/offloaded.
+        if self.classifier.weight.device != pooled.device:
+            self.classifier = self.classifier.to(pooled.device)
+
         # ── 3. Classify ──
-        logit = self.classifier(pooled)    # (B, 1)
+        logit = self.classifier(pooled.float())    # (B, 1)
+        logit = torch.nan_to_num(logit, nan=0.0, posinf=20.0, neginf=-20.0)
 
         result = {"logit": logit}