device_map can not use xla device such as tpu

[chaowenguo]

Describe the bug

model_id = 'Wan-AI/Wan2.1-T2V-14B-Diffusers'
model_dir = huggingface_hub.snapshot_download(model_id)
pipe = diffusers.WanPipeline.from_pretrained(model_dir, vae=diffusers.AutoencoderKLWan.from_pretrained(model_dir, subfolder='vae', torch_dtype=torch.float32), torch_dtype=torch.bfloat16, device_map='balanced')
pipe.scheduler = diffusers.UniPCMultistepScheduler(prediction_type='flow_prediction', use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=5)
print(pipe.hf_device_map)
prompt = 'a beautiful japanese girl'
negative_prompt = 'ugly, monochrome'
output = pipe(prompt=prompt, negative_prompt=negative_prompt, height=1280, width=720, num_frames=17, guidance_scale=5.0).frames[0]
diffusers.utils.export_to_video(output, 'output.mp4', fps=8)

I try to run the above code in kaggle tpu, but I do not see anything output in print(pipe.hf_device_map) and I just see cpu usage become huge, but not tpu usage. How to use device_map with tpu?

Reproduction

N/A

Logs

System Info

Diffusers 0.33.1

Who can help?

No response

[a-r-r-o-w]

cc @SunMarc

[chaowenguo]

@SunMarc any fee back of my question？