Dictionary model inputs .to(device) issue

[mshuaibi]

Ahhh that makes sense now. Thank you!