[Deep Learning] [Pytorch] Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! 에러

 

 

내용 출처 : https://sweetdev.tistory.com/946

[ML][PyTorch] Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! 에러 해결하기

 

이유 : 

 

check list 

* 데이터가 텐서로 변환이 되었는가?

* 데이터끼리 data type 이 같은가 

* 모델 입력되기전 데이터가 모델과 같은 device 인가

* 모델들은 다 같은 device로 통일 되었는가

 

 

.to_device("cuda") 를 하지 않으면 모든 연산은 cpu에서 작동한다. 

처음부터 gpu로 시작했으면 모든 데이터, 텐서연산, layer에 모두 gpu device를 쓴다고 지정해줘야한다. 

귀찮지만 그래야한다. 

 

 gpu device를 쓴다고 지정하는 방법은 다양하지만 나는 아래와 같이 하기로 했따 ! 

class LinearLayer(nn.Module):
    def __init__(self, input_dim, output_dim, device):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.linear = nn.Linear(input_dim, output_dim, bias = True).to(device) ### gpu
    
    def forward(self, x):
        x_flatten = x.view(-1, self.input_dim).to(device) ### gpu
        return self.linear(x_flatten)
        
def model_setting(input_dim, output_dim, device):

    criterion = nn.CrossEntropyLoss().to(device) # 내부적으로 소프트맥스 함수를 포함하고 있음. ## gpu
    model = LinearLayer(input_dim, output_dim, device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    return criterion, model, optimizer
    
def training(training_epochs, data_loader, criterion, model, optimizer, device):
    for one_epoch in range(training_epochs):
        total_batch, avg_cost = len(data_loader), 0
        for x_train, y_train in data_loader:
            y_train = y_train.to(device) #### gpu
            model_output = model(x_train)
            cost = criterion(model_output, y_train) ## CrossEntropyLoss

            cost.backward()
            optimizer.step()
            optimizer.zero_grad()

            avg_cost += cost / total_batch

        print(f"Epoch : {one_epoch+1} / {training_epochs}, cost : {avg_cost:.4f}")
    return model