From a4f0b57ae856c732d65435181bbc63edf6286537 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Mon, 4 Sep 2023 01:27:16 +0800
Subject: [PATCH] feat: add mps to chinese bert

---
 text/chinese_bert.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/text/chinese_bert.py b/text/chinese_bert.py
index cb84ce0..45e5e28 100644
--- a/text/chinese_bert.py
+++ b/text/chinese_bert.py
@@ -1,7 +1,16 @@
 import torch
+import sys
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device(
+        "cuda"
+        if torch.cuda.is_available()
+        else (
+            "mps"
+            if sys.platform == "darwin" and torch.backends.mps.is_available()
+            else "cpu"
+        )
+    )
 
 tokenizer = AutoTokenizer.from_pretrained("./bert/chinese-roberta-wwm-ext-large")
 model = AutoModelForMaskedLM.from_pretrained("./bert/chinese-roberta-wwm-ext-large").to(device)