Update app.py

app.py

@@ -80,7 +80,7 @@ def embed(req: EmbedRequest):
         # Case 1: Query → mean pool across token embeddings
         # -----------------------------
         if (req.prompt_name or "").lower() == "query":
-            with torch.inference_mode():
+            with torch.no_grad():
                 outputs = model.encode_text(
                     texts=[text],
                     task=req.task,
@@ -107,7 +107,7 @@ def embed(req: EmbedRequest):
             end = min(position + max_len, total_tokens)
             window_ids = input_ids[position:end].unsqueeze(0).to(device)
     
-            with torch.inference_mode():
+            with torch.no_grad():
                 outputs = model.encode_text(
                     texts=[tokenizer.decode(window_ids[0])],
                     task=req.task,
@@ -139,7 +139,7 @@ def embed(req: EmbedRequest):
 @app.post("/embed_image", response_model=EmbedImageResponse)
 def embed_image(req: EmbedImageRequest):
     try:
-        with torch.inference_mode():
+        with torch.no_grad():
             outputs = model.encode_image(
                 images=[req.image],
                 task=req.task,