Clamp logits BEFORE softmax

app.py

@@ -135,6 +135,7 @@ def generate_diffusion_text(input_ids):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
         logits = model(input_ids=input_tensor)["logits"]
+        logits = logits.clamp(min=-1e4, max=1e4)
         probs = torch.nn.functional.softmax(logits, dim=-1)[0]
         probs = torch.clamp(probs, min=1e-8, max=1.0)
         print("probs", probs)