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():
device = 0
elif torch.backends.mps.is_available():
device = "mps"
device = 'mps'
else:
device = -1
self.pipe = pipeline('text-classification', model=model_ckpt, tokenizer=tokenizer, device=device)