Bug: Input device hardcode need to be adjusted

[JeffWang987]

Hi, dear OpenAI contributors.

I encountered an issue while running the following code:

self.consistency_dec = ConsistencyDecoder(device=model.device)  # maybe cuda:3
self.consistency_dec(latents / self.vae.config.scaling_factor)

The problem likely arises from a hardcoding that forces the input set to be on cuda:0, which doesn't match the device of consistency_dec.

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/dalle_runner_api/model_infra/modules/public_diff_vae.py", line 91, in forward
    **_input = torch.to(features, torch.device("cuda:0"), 6)_**
    features0 = torch.upsample_nearest2d(input, None, [8., 8.])