Feat: CPU inference support

2023-12-30 19:52:01 +09:00
parent 4108478f54
commit 9045fa3426
5 changed files with 77 additions and 0 deletions
--- a/text/chinese_bert.py
+++ b/text/chinese_bert.py
@@ -27,6 +27,8 @@ def get_bert_feature(
        device = "mps"
    if not device:
        device = "cuda"
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
    if device not in models.keys():
        models[device] = AutoModelForMaskedLM.from_pretrained(LOCAL_PATH).to(device)
    with torch.no_grad():