Update ruaccent/omograph_model.py

ruaccent/omograph_model.py

@@ -2,8 +2,8 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 
 class OmographModel:
-    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.nli_model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16).to(self.device)