Update app.py

app.py

@@ -83,10 +83,9 @@ def predict(content, policy):
     input_text = PROMPT.format(policy=policy, content=content)
     input_ids = tokenizer.encode(input_text, return_tensors="pt")
 
-    with torch.no_grad():
+    with torch.inference_mode():
         outputs = model(input_ids)
         logits = outputs.logits[:, -1, :]  # Get logits for the last token
-        model.gradient_checkpointing_enable()
         predicted_token_id = torch.argmax(logits, dim=-1).item()
     decoded_output = tokenizer.decode([predicted_token_id])
     if decoded_output == '0':