nn.module에서 CPU, GPU 혼합 사용으로 인한 error

HeyHo·2022년 11월 16일
0
#all tensors of each operation should be in same device
class Module(nn.Module):
  def __init__(self):
    super().__init__()
    self.network = nn.Sequential(nn.Linear(1000, 1000), nn.Linear(1000, 100))

  def forward(self, x):
    return self.network(x)

다음과 같이 module이 정의될 때, 아래와 같이 연산을 진행해보았다.

module = Module().to(0)
x = torch.zeros(1000, 1000).to(0)
module(x)


module에서 to(0)연산을 통해 module의 parameter를 GPU로 보냈지만, x는 여전히 cpu에 남아있기 때문에
'RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!' 에러가 발생하였다.

module = Module().to(0)
x = torch.zeros(1000, 1000)
module(x)

똑같이 x의 device를 GPU로 보내주면, device에서 없이 연산이 잘 된 것을 확인할 수 있다.

profile
Coputer vision, AI

0개의 댓글