update cuda.

shibing624 · Jul 31, 2024 · 0a643f1 · 0a643f1
1 parent 32e5c46
commit 0a643f1
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/pytextclassifier/bert_classification_model.py b/pytextclassifier/bert_classification_model.py
@@ -319,8 +319,7 @@ def __init__(
                     " Make sure CUDA is available or set use_cuda=False."
                 )
         else:
-            self.device = 'cuda' if torch.cuda.is_available() else (
-                'mps' if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else 'cpu')
+            self.device = 'mps' if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() else 'cpu'
         logger.debug(f"Device: {self.device}")
 
         self.loss_fct = init_loss(