Skip to content

[已解決] RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cpu and parameter tensor at cuda:0

今天,我在使用 PyTorch 讀取一個已經訓練好的模型進行全新資料的分類時,意外地遇到了以下這個報錯:

RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cuda:1 and parameter tensor at cuda:0

這個報錯其實與我標題下得不太一樣。這是因為在我上網查詢之後,我發現這個報錯較多發生在讀取『已經訓練好的模型』後、接著再使用 CPU 驅動模型進行分類的情況。

為什麼會發生這個報錯呢?其實報錯中也都說明了:

Input and parameter tensors are not at the same device.

很有可能,我們讀取的模型預設是放在 GPU 的記憶體中、而我們要用來處理的新資料,則是放在 CPU 中。


解決方法

解決方法其實說穿了異常單純:既然模型預設是放在 GPU 的記憶體中 (畢竟當初就是在 GPU 上訓練的),那麼,我們只要將模型一起放到 CPU 中即可。

大致上可能的程式如下:

import torch

# Model
model = torch.load('MODEL_NAME').cpu()



這樣一來,模型就從原先預設的 GPU 移動至 CPU 了。

順帶一提,我真實遇見的情況是我模型讀取進來預設是放在第一顆 GPU、然而我想要重新處理的資料放在第二顆 GPU。

所以,在我的案例中,我將模型移動至了第二顆 GPU。

import torch

# Model
model = torch.load('MODEL_NAME').to('cuda:1')



就像這樣。


References

Leave a Reply