Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix CUDA/CPU device mismatch in quality/eval.py #142

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

politecat314
Copy link

When running quality/eval.py with device="cuda", the following error occurred on line 43:

pred_score = net(input_data).data.cpu().numpy().squeeze()
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

This error occurred because the model was on cuda but the input_data was on cpu.

Changes made to resolve the issue:

  1. Modified read_img function to accept a device parameter and move the input tensor to the specified device.
  2. Updated the device selection in the main script to use torch.device() with a check for CUDA availability.
  3. Ensured that the input data is moved to the same device as the model before forward pass.

These changes ensure that both the model and input data are on the same device, resolving the CUDA/CPU mismatch error.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant