support cpu training, use cpu training on mac

This commit is contained in:
Lion
2024-03-13 22:09:24 +08:00
parent 931781774d
commit 1963eb01cc
8 changed files with 41 additions and 43 deletions

View File

@@ -49,8 +49,8 @@ if os.path.exists(txt_path) == False:
os.makedirs(bert_dir, exist_ok=True)
if torch.cuda.is_available():
device = "cuda:0"
elif torch.backends.mps.is_available():
device = "mps"
# elif torch.backends.mps.is_available():
# device = "mps"
else:
device = "cpu"
tokenizer = AutoTokenizer.from_pretrained(bert_pretrained_dir)

View File

@@ -50,8 +50,8 @@ maxx=0.95
alpha=0.5
if torch.cuda.is_available():
device = "cuda:0"
elif torch.backends.mps.is_available():
device = "mps"
# elif torch.backends.mps.is_available():
# device = "mps"
else:
device = "cpu"
model=cnhubert.get_model()

View File

@@ -40,8 +40,8 @@ if os.path.exists(semantic_path) == False:
if torch.cuda.is_available():
device = "cuda"
elif torch.backends.mps.is_available():
device = "mps"
# elif torch.backends.mps.is_available():
# device = "mps"
else:
device = "cpu"
hps = utils.get_hparams_from_file(s2config_path)