Fix torch device issues (#20304)

Commit

3 years ago

Fix torch device issues (#20304) * fix device issue Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>