在 pytorch mps 设备 M1 pro 上使用 HuggingFace 管道

Question

我想在 mps 设备上运行零样本分类任务的管道摘要。这是我的代码

pipe = pipeline('zero-shot-classification', device = mps_device)
seq = "i love watching the office show"
labels = ['negative', 'positive']
pipe(seq, labels)

产生的错误是

RuntimeError: Placeholder storage has not been allocated on MPS device!

我的猜测是因为 seq 在我的 cpu 上而不是 mps 上。我怎样才能解决这个问题 ？有没有办法将 seq 发送到 mps 设备，以便我可以将其传递到管道进行推理？

谢谢

Answer 1

当我遇到类似的问题时，它通过这样做得到了解决model = model.to("mps")，尽管这在你的情况下不应该是一个问题。

以下代码在我的机器上运行：

import os
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

from transformers import pipeline
mps_device = "mps"
pipe = pipeline('zero-shot-classification', device = mps_device)
seq = "i love watching the office show"
labels = ['negative', 'positive']
pipe(seq, labels)