Pytorch Albert抱着脸，同时在CPU和GPU上加载模型

当试图在GPU上加载模型时，模型也加载到CPU中 下面的代码以两种方式加载模型

MODEL_DEVICE = torch.device('cuda')
MODEL_PATH='./models'
tokenizer = AlbertTokenizer.from_pretrained(MODEL_PATH)
qa_model = AlbertForQuestionAnswering.from_pretrained(MODEL_PATH).to(MODEL_DEVICE)

我已经参考了下面的链接以供参考。

该型号同时占用RAM内存（1.5 GB）和GPU内存（650 MB）。 我已经将torch设备指定为Cuda，但它的行为仍不符合预期。 当指定“cpu”时，它工作良好，不会加载到GPU中。但当指定cuda时，它也会在CPU和GPU中加载。 我试过，“cuda”，“cuda:0”

我还尝试将拥抱脸库代码从

state_dict = torch.load(resolved_archive_file, map_location="cpu")

到

---

有什么解决方案吗？

通过迭代state_dict（）.items（）并检查类型，确认每个张量的类型，可以获得有关模型的线索。

state_dict = torch.load(resolved_archive_file, map_location="cuda")