Add device and key for lazy tensors (#61621)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/61621
Test Plan: CI
Reviewed By: mruberry
Differential Revision: D29912934
Pulled By: asuhan
fbshipit-source-id: 493c32063a3e756d93cbf1d876563a35eaafb537
