Use autocast to float16

app.py

@@ -77,7 +77,8 @@ def confidence_guided_noising(input_ids, answer_start, confidences, noise_clippi
 def generate_diffusion_text(input_ids, top_p, top_k):
     with torch.no_grad():
         input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device)
-        logits = model(input_ids=input_tensor)["logits"]
+        with torch.amp.autocast('cuda', dtype=torch.float16):
+            logits = model(input_ids=input_tensor)["logits"]
         logits = filter_logits(logits, top_k=top_p, top_p=top_k) 
         logits = logits.clamp(min=-1e8, max=1e4)
         probs = torch.nn.functional.softmax(logits, dim=-1)[0]