Fix complex rewrite device mismatch

[jloftin-nv]

Description

After complex-graph rewrite, FakeTensorProp can mix CUDA activations with CPU _frozen_param* weights and raise FakeTensorDeviceMismatchError, aborting builds that dry-ran cleanly.

Align parameter devices before prop (and keep the non-tensor meta["val"] guard) so rewritten graphs re-propagate on a single device.

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• [ X] My code follows the style guidelines of this project (You can use the linters)
• [ X] I have performed a self-review of my own code
• [ X] I have commented my code, particularly in hard-to-understand areas and hacks
• [ X] I have made corresponding changes to the documentation
• [ X] I have added tests to verify my fix or my feature
• [ X] New and existing unit tests pass locally with my changes
• [ X] I have added the relevant labels to my PR in so that relevant reviewers are notified

[apbose]

this test would pass even without the fix. 0 dim would pass in eager mode and fake tensor mode,but then if you do torch.tensor([2.0]) triggers a mismatch too early FakeTensorDeviceMismatchError—in eager/export—rather than specifically during post-rewrite metadata propagation.

[apbose]

Every tensor attribute is moved to the first CUDA input’s device.
The new fetch_attr changes all tensor attributes whenever any CUDA placeholder exists. That includes intentionally CPU attributes or an independent CPU branch, producing incorrect device metadata. Something like below

 class M(torch.nn.Module):
      def __init__(self):
          super().__init__()
          self.register_buffer("cpu_table", torch.arange(4))  # intentionally CPU

      def forward(self, z):
          gpu_result = z * z                 # z is CUDA complex
          cpu_result = self.cpu_table + 1    # independent CPU computation
          return gpu_result, cpu_result

Maybe we should add

• A real _frozen_param mismatch test proving the target attribute gets aligned.
• A mixed-device test like the example above proving an unrelated CPU attribute remains CPU.

[apbose]

LGTM. Waiting on CI.