Repository: pytorch
Issue: Diffuser pipeline device attribute broken when using optimized model
Labels: triaged,bug

### 🐛 Describe the bug

Hi,

just adding this for tracking purposes as a solution is potentially provided by [this PR](https://github.com/pytorch/pytorch/pull/88164).

The issue is that when using a compiled model in a [DDPMPipeline](https://github.com/huggingface/diffusers/blob/v0.6.0/src/diffusers/pipelines/ddpm/pipeline_ddpm.py#L24) the self.device property does not give the correct result. It [always return "cpu"](https://github.com/huggingface/diffusers/blob/ad9d7ce4763f8fb2a9e620bff017830c26086c36/src/diffusers/pipeline_utils.py#L213) as the optimized model class does not inherent from nn.Module.

Here is a simple repro:
```
import torch
import torch._dynamo as dynamo
from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel

model = UNet2DModel(
    sample_size=128,  # the target image resolution
    in_channels=3,  # the number of input channels, 3 for RGB images
    out_channels=3,  # the number of output channels
    layers_per_block=2,  # how many ResNet layers to use per UNet block
    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channes for each UNet block
    down_block_types=( 
        "DownBlock2D",  # a regular ResNet downsampling block
        "DownBlock2D", 
        "DownBlock2D", 
        "DownBlock2D", 
        "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
        "DownBlock2D",
    ), 
    up_block_types=(
        "UpBlock2D",  # a regular ResNet upsampling block
        "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
        "UpBlock2D", 
        "UpBlock2D", 
        "UpBlock2D", 
        "UpBlock2D"  
      ),
)

if torch.cuda.is_available():
    model.to('cuda')

model = torch._dynamo.optimize('inductor')(model)

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

pipeline = DDPMPipeline(unet=model, scheduler=noise_scheduler)

if torch.cuda.is_available():
    pipeline.to('cuda')

pipeline(
    batch_size = 16, 
    generator=torch.manual_seed(42),
    )
```

PR [88164](https://github.com/pytorch/pytorch/pull/88164) will solve this as it makes a compiled model a nn.Module.

### Error logs

```
Traceback (most recent call last):
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/tensor.py", line 131, in _get_fake_value
    return wrap_fake_exception(
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 725, in wrap_fake_exception
    return fn()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/tensor.py", line 132, in <lambda>
    lambda: _run_node(tx.output, node, args, kwargs, nnmodule)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/tensor.py", line 56, in _run_node
    return nnmodule(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1423, in _call_impl
    return forward_call(*input, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 860, in __torch_dispatch__
    return decomposition_table[func](*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_prims_common/wrappers.py", line 212, in _fn
    result = fn(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_decomp/decompositions.py", line 63, in inner
    r = f(*tree_map(increase_prec, args), **tree_map(increase_prec, kwargs))
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_decomp/decompositions.py", line 1111, in addmm
    out = alpha * torch.mm(mat1, mat2)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 901, in __torch_dispatch__
    return self.wrap_meta_outputs_with_default_device_logic(r, func, args, kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 934, in wrap_meta_outputs_with_default_device_logic
    return tree_map(partial(wrap), r)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/utils/_pytree.py", line 195, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/utils/_pytree.py", line 195, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 951, in wrap
    ) = FakeTensor._find_common_device(func, args, kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 650, in _find_common_device
    tree_map(merge_devices, args)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/utils/_pytree.py", line 195, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/utils/_pytree.py", line 195, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_subclasses/fake_tensor.py", line 646, in merge_devices
    raise RuntimeError(
RuntimeError: Unhandled FakeTensor Device Propagation for aten.mm.default, found two different devices cpu, cuda:0

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/home/mreso/torchdynamo/diffuser_training/train_repro.py", line 46, in <module>
    pipeline(
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/diffusers/pipelines/ddpm/pipeline_ddpm.py", line 84, in __call__
    model_output = self.unet(image, t).sample
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 137, in __call__
    return self.forward(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 134, in forward
    return optimized_forward(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 236, in catch_errors
    return callback(frame, cache_size)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 466, in _convert_frame
    result = inner_convert(frame, cache_size)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 118, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 92, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 339, in _convert_frame_assert
    return _compile(
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 394, in _compile
    out_code = transform_code_object(code, transform)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py", line 341, in transform_code_object
    transformations(instructions, code_options)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 382, in transform
    tracer.run()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1452, in run
    super().run()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 352, in run
    and self.step()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 322, in step
    getattr(self, inst.opname)(inst)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 174, in wrapper
    return inner_fn(self, inst)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 758, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 264, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/nn_module.py", line 221, in call_function
    return tx.inline_user_function_return(
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 293, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1524, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 1578, in inline_call_
    tracer.run()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 352, in run
    and self.step()
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 322, in step
    getattr(self, inst.opname)(inst)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 174, in wrapper
    return inner_fn(self, inst)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 758, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 264, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/nn_module.py", line 201, in call_function
    return variables.TensorVariable.create(
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/tensor.py", line 201, in create
    example_value = _get_fake_value(proxy.node, tx)
  File "/fsx/users/mreso/conda/envs/tdn/lib/python3.9/site-packages/torch/_dynamo/variables/tensor.py", line 145, in _get_fake_value
    raise TorchRuntimeError() from e
torch._dynamo.exc.TorchRuntimeError:

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torchdynamo.config.suppress_errors = True
==========
```

### Minified repro

Minifier does not work here but I am not sure if it should in these cases.