Gracefully handle extra (non-tensor) items in transfer_batch_to_device #20914
robmarkcole
started this conversation in
Idea pool
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
-
Problem
When a Batch contains non-tensor items (e.g. a str such as an
image_id
or other metadata) Lightning tries to call .to(device) on it and crashes:TypeError: can't move str object to device cuda:0
Where this surfaces
It is useful to add keys such as "metadata" to each sample so that callbacks or logging utilities can trace predictions back to the original image. These keys often hold strings or other Python objects that do not need to be on the GPU.
Impact
Users must strip such metadata before training and manually re-attach it afterwards, adding boilerplate and risk of error. It also discourages passing useful context through the training loop.
CC @adamjstewart
Beta Was this translation helpful? Give feedback.
All reactions