Update ruaccent/accent_model.py

ruaccent/accent_model.py

@@ -3,8 +3,8 @@ from .char_tokenizer import CharTokenizer
 from transformers import AutoModelForTokenClassification
 
 class AccentModel:
-    def __init__(self) -> None:
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    def __init__(self, allow_cuda=True) -> None:
+        self.device = torch.device('cuda' if torch.cuda.is_available()  and allow_cuda else 'cpu')
     def load(self, path):
         self.model = AutoModelForTokenClassification.from_pretrained(path).to(self.device)
         self.tokenizer = CharTokenizer.from_pretrained(path)