[에러 해결 / pytorch] RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

개발환경

 

- windows 10

- python 3.9.7

- pytorch

 

1. Problem

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

Input type과 weight type 이 같아야 한다는 Runtime Error

 

2. Reason

GPU 학습 시, model이나 dataset을 GPU 에 할당하지 않았기 때문에 발생합니다.

 

 

3. Solution

model을 GPU에 올리는 코드

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)

 

 

dataset을 GPU에 올리는 예시 코드

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_set = train_set.to(device)
val_set = val_set.to(device)
test_set = test_set.to(device)

 

Training Loop 안에서 input과 labels를 같은 GPU에 올려주어야 합니다 !

아래는 예시 코드 일부입니다.

# Training loop
for epoch in range(num_epochs):
    for batch in train_loader:
        # Move batch data to GPU
        inputs = batch[0].to(device)
        labels = batch[1].to(device)

        # Perform forward and backward pass
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

 

도움 되었다면 좋아요 부탁드려요 :)