From 96c0a81e71b5817d204bc4b9c9bf35c7b0d49678 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= Date: Tue, 5 Sep 2023 10:24:25 +0800 Subject: [PATCH] support mps --- text/japanese_bert.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/text/japanese_bert.py b/text/japanese_bert.py index 304d03a..bb83696 100644 --- a/text/japanese_bert.py +++ b/text/japanese_bert.py @@ -1,7 +1,15 @@ import torch 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/bert-base-japanese-v3") model = AutoModelForMaskedLM.from_pretrained("./bert/bert-base-japanese-v3").to(device)