如何在CUDA上用PyTorch运行CIFAR-10？GPU运行报错求助

Hey there! Let's get your code running on GPU smoothly. The error you're seeing usually happens because your model and input tensors aren't on the same device (CPU vs GPU), which causes a mismatch during computation. Here's how to fix it step by step, tailored to your PyTorch 0.4.0 setup:

1. Set Up Your Device First #

Start by defining a device variable to handle GPU/CPU switching gracefully. This is cleaner than calling .cuda() everywhere:

import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

First, confirm this prints cuda:0—if it doesn't, double-check your PyTorch GPU installation (even though TensorFlow works, PyTorch might not have been built with CUDA support; but since you're using 0.4.0 with CUDA 9.0, this should be fine).

2. Move Your Model to the GPU #

Once you've defined your Net class, move the entire model to the GPU using .to(device) (or .cuda() if you prefer, but .to() is more flexible):

net = Net().to(device)

This ensures all layers (like conv1, pool, etc.) are on the GPU, so their operations will run there by default.

3. Move Input Tensors to the GPU #

In your training loop, before passing data to the model, move both the images and labels to the same device as the model:

for i, data in enumerate(trainloader, 0):
    inputs, labels = data
    # Shift data to GPU
    inputs, labels = inputs.to(device), labels.to(device)
    
    # Rest of your training code (zero grad, forward, backward, optimize)
    optimizer.zero_grad()
    outputs = net(inputs)  # Now this runs entirely on GPU
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

Why Your Original Code Failed #

Looking at your error trace, you added .cuda() inside the forward method:

x = self.pool(F.relu(self.conv1(x)).cuda())

This is problematic because:

• If your model is on CPU but you try to move the tensor to GPU mid-forward pass, you get a device mismatch.
• Even if the model is on GPU, manually adding .cuda() here is redundant and can cause issues if inputs aren't already on GPU.

Remove that .cuda() call from your forward method—once the model and inputs are on GPU, all operations in forward will automatically run on GPU.

Quick Validation #

After making these changes, you can verify tensors are on GPU by checking:

print(inputs.device)  # Should show cuda:0
print(next(net.parameters()).device)  # Should also show cuda:0

If both match, you're good to go!

内容的提问来源于stack exchange，提问作者Kalpit