nightly robustness fixes for linking across devices (#43771)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/43761
CC rgommers ezyang
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43771
Reviewed By: glaringlee
Differential Revision: D23819835
Pulled By: malfet
fbshipit-source-id: a3be2780c4b8bdbf347d456c4d14df863c2ff8c2
