Workaround for hardcoded device in PT2 IR (#2098)

Summary:
Pull Request resolved: #2098

PT2 IR can end up hardcoding the device during tracing, for example: https://fburl.com/code/m478zn5s. This diff provides a workaround util function move_to_copy_nodes_to_device to correctly set hardcoded devices to an appropriate file.

Reviewed By: wilson100hong

Differential Revision: D58426261

fbshipit-source-id: 9435661b03a4853558059e243392896f519c4565

torchrec/ir/utils.py

@@ -151,3 +151,22 @@ def mark_dynamic_kjt(
         if kjt._offsets is not None:
             shapes_collection[kjt._offsets] = (olen,)
     return shapes_collection
+
+
+def move_to_copy_nodes_to_device(
+    unflattened_module: nn.Module,
+    device: torch.device,
+) -> nn.Module:
+    """
+    Moves all the copy nodes to the given device.
+    """
+    for nodes in unflattened_module.graph.nodes:
+        if "_to_copy" in nodes.name:
+            new_kwargs = {}
+            for k, v in nodes.kwargs.items():
+                if isinstance(v, torch.device):
+                    v = device
+                new_kwargs[k] = v
+            nodes.kwargs = new_kwargs
+
+    return unflattened_module