Fix a device inconsistency in dlrm (#1170)
Summary:
See the latest problem description in https://github.com/pytorch/torchdynamo/issues/1023
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1170
Reviewed By: xuzhao9
Differential Revision: D39366515
Pulled By: desertfire
fbshipit-source-id: 733a905d286c64e3e6bbc229bef2be2e3806d8c3
