FlowNet2-C推理时输入与权重设备类型不匹配问题求助

That error happens because your model's weights are running on the CPU while your input data is being processed on the GPU (or vice versa)—they need to live on the same device to work together. Here's how to fix this step by step:

1. Set up your target device correctly #

First, confirm if CUDA is available and define your device:

import torch

# Use GPU if available, otherwise fall back to CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2. Load and move your model to the device #

When loading your pre-trained FlowNet2-C model, make sure to shift the entire model to your target device. Here's how:

# Import the model class (adjust the path based on the repo's structure)
from models import FlowNet2C

# Initialize the model architecture
model = FlowNet2C()

# Load the pre-trained weights
# Option 1: Load first, then move to device
model.load_state_dict(torch.load("path/to/your/FlowNet2-C.pth"))
model = model.to(device)

# Option 2: Load directly to the device (handy if the model was saved on GPU)
# model.load_state_dict(torch.load("path/to/your/FlowNet2-C.pth", map_location=device))

# Switch to evaluation mode (disables training-only layers like dropout)
model.eval()

3. Move your input data to the same device #

Before feeding your input (like image pairs) into the model, ensure it's on the same device as the model:

# Example: Prepare your preprocessed input tensor (adjust to your data pipeline)
input_data = ...  # Your image pair tensor ready for inference

# Move input to the target device
input_data = input_data.to(device)

# Run inference without gradient computation (faster and saves memory)
with torch.no_grad():
    flow_output = model(input_data)

Common Mistakes to Watch For #

• Forgetting to call model.to(device) after loading weights—this leaves the model stuck on CPU even if your input is on GPU.
• Skipping the model.eval() call—this can lead to unexpected behavior from training-specific layers during inference.
• Neglecting to move input tensors to the device—double-check your data loading code to ensure inputs aren't staying on CPU.

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