fix: move predict_image input tensor to the model's device

[giswqs]

Summary

Fixes #1390.

deepforest.predict._predict_image_ builds the input tensor on CPU and never moves it to the model's device, so any caller that has placed the model on CUDA (model.model.to('cuda'), or a Lightning trainer that has trained on GPU) gets:

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 PR moves the input tensor to next(model.parameters()).device before the forward pass. Guarded with try/except StopIteration to handle the (unrealistic) case of a parameter-less model gracefully.

Repro

from deepforest import main
m = main.deepforest()
m.load_model(model_name='weecology/deepforest-tree', revision='main')
m.model.to('cuda')
m.predict_image(path='some_rgb.tif')
# before this PR: RuntimeError as above
# after this PR: returns a DataFrame of detections

Changes

• src/deepforest/predict.py: move the input tensor to the model's device inside _predict_image_.
• tests/test_main.py: add test_predict_image_on_cuda regression test, skipped when CUDA is unavailable.

Test plan

• ruff format --check and ruff check on the changed source files (clean, using the version pinned in .pre-commit-config.yaml).
• CI: existing CPU predict_image tests should still pass; new CUDA test is auto-skipped on CPU-only runners.
• Verified locally that the workaround applied downstream in opengeos/geoai#746 becomes unnecessary once this lands.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.84%. Comparing base (f3bd776) to head (9006f50).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
src/deepforest/predict.py	50.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1391      +/-   ##
==========================================
- Coverage   86.96%   86.84%   -0.12%     
==========================================
  Files          26       26              
  Lines        3712     3717       +5     
==========================================
  Hits         3228     3228              
- Misses        484      489       +5     

Flag	Coverage Δ
unittests	86.84% <50.00%> (-0.12%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[bw4sz]

Thanks @giswqs! @jveitchmichaelis will have a look at this.

[jveitchmichaelis]

Thanks, this looks pretty sensible to me. I'll double check I can reproduce too.

[jveitchmichaelis]

@giswqs thanks for raising this, I can replicate although I was surprised that it failed. The additional cause is that the model isn't moved to the device until a trainer uses it and predict_image is a low level call that doesn't invoke trainer.predict which would normally handle device placement.

Alternatively, you could use tiled prediction and I think you would avoid this (if you want a single tile, should work if you set the patch size to the size of the image).

So I'd propose:

1. I think this solution is fine in the short term, if this is how you're using the model downstream.
2. Longer term I think we would replace the guts of this function with a call to trainer.predict using a SingleImage dataset (or MultiImage). However this would only work with paths, so maybe (1) is still required if the user passes an array.

I'll leave some comments on the test case in the review.

@bw4sz any thoughts there? On the one hand I feel we shouldn't encourage manually moving the model, because accelerator is in the config. On the other, if you specify cuda, the model doesn't actually get moved until the trainer is invoked. The complication is that Lightning also handles device names for us, so if we tried to add that in predict_image we would need to add an auto -> expected device helper.

[bw4sz]

We can allow this for now, but definitely let's get predict_image into the trainer context that feels like the right pattern.