This commit is contained in:
parent
ca8f4a43bf
commit
0411000c03
10
app/main.py
10
app/main.py
|
|
@ -185,7 +185,15 @@ def _encode(texts: List[str]):
|
||||||
"""
|
"""
|
||||||
def _worker(t, q):
|
def _worker(t, q):
|
||||||
try:
|
try:
|
||||||
q.put(("ok", model.encode(t, return_dense=True)))
|
if DEVICE.startswith("cuda"):
|
||||||
|
# 正常 GPU 跑,多进程池照旧
|
||||||
|
out = model.encode(t, return_dense=True)
|
||||||
|
else:
|
||||||
|
# 已经 fallback 到 CPU:禁用进程池,禁止再向 GPU 拷权重
|
||||||
|
out = model.encode(t, return_dense=True,
|
||||||
|
num_processes=1,
|
||||||
|
device="cpu")
|
||||||
|
q.put(("ok", out))
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
q.put(("err", str(e)))
|
q.put(("err", str(e)))
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue