Update app.py

app.py

@@ -26,15 +26,26 @@ def generate(prompt, max_new_tokens, temperature, top_k):
     with torch.no_grad():
         for _ in range(max_new_tokens):
             outputs = model(input_ids)
-            logits = outputs["logits"][:, -1, :]
 
-            logits = logits / max(temperature, 1e-5)
+            # Cast to fp32 for the sampling math — fp16 logits from an
+            # early/undertrained checkpoint can overflow fp16's range,
+            # which turns softmax output into inf/nan and crashes
+            # multinomial with a CUDA device-side assert.
+            logits = outputs["logits"][:, -1, :].float() / max(temperature, 1e-5)
 
             if top_k > 0:
                 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                 logits[logits < v[:, [-1]]] = -float("Inf")
 
             probs = torch.softmax(logits, dim=-1)
+
+            # Defensive fallback: if anything still slips through as
+            # nan/inf/negative, fall back to a safe uniform distribution
+            # instead of crashing the whole generation call.
+            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+            if probs.sum().item() <= 0:
+                probs = torch.ones_like(probs) / probs.size(-1)
+
             next_token = torch.multinomial(probs, num_samples=1)
 
             input_ids = torch.cat([input_ids, next_token], dim=-1)