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

Last Updated on 2021-06-28 by Clay

今天，我在使用 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()

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')

import torch

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

就像這樣。