Преглед на файлове

fix: support NPU device in UnimernetModel initialization

myhloli преди 5 месеца
родител
ревизия
2785f60424
променени са 1 файла, в които са добавени 1 реда и са изтрити 1 реда
  1. 1 1
      mineru/model/mfr/unimernet/Unimernet.py

+ 1 - 1
mineru/model/mfr/unimernet/Unimernet.py

@@ -21,7 +21,7 @@ class MathDataset(Dataset):
 class UnimernetModel(object):
     def __init__(self, weight_dir, _device_="cpu"):
         from .unimernet_hf import UnimernetModel
-        if _device_.startswith("mps"):
+        if _device_.startswith("mps") or _device_.startswith("npu"):
             self.model = UnimernetModel.from_pretrained(weight_dir, attn_implementation="eager")
         else:
             self.model = UnimernetModel.from_pretrained(weight_dir)