Модель pytorch, сохраненная из TPU, запускается на CPU

Нашел интересную модель - генератор вопросов, но не могу запустить. Я получил ошибку:

Traceback (most recent call last):
  File "qg.py", line 5, in <module>
    model = AutoModelWithLMHead.from_pretrained("/home/user/ml-experiments/gamesgen/t5-base-finetuned-question-generation-ap/")
  File "/home/user/.virtualenvs/hugging/lib/python3.7/site-packages/transformers/modeling_auto.py", line 806, in from_pretrained
    return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
  File "/home/user/.virtualenvs/hugging/lib/python3.7/site-packages/transformers/modeling_utils.py", line 798, in from_pretrained
    import torch_xla.core.xla_model as xm
ModuleNotFoundError: No module named 'torch_xla'

Я быстро погуглил и обнаружил, что torch_xla — это то, что используется для обучения модели pytorch на TPU. Но я хотел бы запустить его локально на процессоре (конечно, для вывода) и получил эту ошибку, когда pytorch попытался загрузить тензоры, привязанные к tpu. Как я могу это исправить?

это модель, которую я пробовал: https://huggingface.co/mrm8488/t5-base-finetuned-question-generation-ap