Fix torch device issue
# What does this PR do?

The fixes in #20304 is changed somehow to the wrong places in #20160, and we got the torch device issues.
This PR fixes this device issue - just put `to` in the correct places.
_The documentation is not available anymore as the PR was closed or merged._