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/english_bert_mock.py
+++ b/text/english_bert_mock.py
@@ -28,6 +28,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] = DebertaV2Model.from_pretrained(LOCAL_PATH).to(device)
    with torch.no_grad():