Only put tensors on a device #5223
Conversation
Codecov Report
@@ Coverage Diff @@
## master #5223 +/- ##
==========================================
- Coverage 77.98% 77.95% -0.04%
==========================================
Files 138 138
Lines 23839 23841 +2
==========================================
- Hits 18592 18586 -6
- Misses 5247 5255 +8
Continue to review full report at Codecov.
|
src/transformers/trainer.py
Outdated
| inputs = {k: v.to(self.args.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} | ||
|
|
There was a problem hiding this comment.
I find this personally slightly hard to read.
Maybe:
for k, v in inputs.items():
if isinstance(v, torch.Tensor):
inputs[k] = v.to(self.args.device)There was a problem hiding this comment.
(otherwise, LGTM – we might want to document this in the function's type signature – or not)
There was a problem hiding this comment.
Usually using list-comprehensions are faster, but the batches don't have a lot of keys so it probably doesn't matter.
src/transformers/trainer.py
Outdated
| inputs = {k: v.to(self.args.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} | ||
|
|
There was a problem hiding this comment.
(otherwise, LGTM – we might want to document this in the function's type signature – or not)
Fix Trainer when users have inputs containing non-tensor values.