This commit is contained in:
AuroraWright
2024-12-12 21:37:14 +01:00
parent 1897324de2
commit 5bf39b4f61

View File

@@ -269,7 +269,7 @@ class TextFiltering:
if torch.cuda.is_available(): if torch.cuda.is_available():
device = 0 device = 0
elif torch.backends.mps.is_available(): elif torch.backends.mps.is_available():
device = "mps" device = 'mps'
else: else:
device = -1 device = -1
self.pipe = pipeline('text-classification', model=model_ckpt, tokenizer=tokenizer, device=device) self.pipe = pipeline('text-classification', model=model_ckpt, tokenizer=tokenizer, device=device)