Dreambooth训练遇RuntimeError：张量设备不匹配（CPU与CUDA）求助

我是一名初学者，在使用Dreambooth训练模型时遇到了设备不匹配的报错，之前查找的解决方案都无效。了解到给变量添加.to(device)可能有用，但不确定该添加到何处，希望有人根据下方报错信息告知添加位置，或提供修复建议。

报错信息：

Traceback (most recent call last):
  File "main.py", line 835, in <module>
    trainer.fit(model, data)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 741, in fit
    self._call_and_handle_interrupt(
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 686, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 778, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1200, in _run
    self._dispatch()
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1280, in _dispatch
    self.training_type_plugin.start_training(self)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\plugins\training_type\training_type_plugin.py", line 202, in start_training
    self._results = trainer.run_stage()
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1290, in run_stage
    return self._run_train()
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1312, in _run_train
    self._run_sanity_check(self.lightning_module)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 1376, in _run_sanity_check
    self._evaluation_loop.run()
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\loops\base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\loops\dataloader\evaluation_loop.py", line 110, in advance
    dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx, dl_max_batches, self.num_dataloaders)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\loops\base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\loops\epoch\evaluation_epoch_loop.py", line 122, in advance
    output = self._evaluation_step(batch, batch_idx, dataloader_idx)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\loops\epoch\evaluation_epoch_loop.py", line 217, in _evaluation_step
    output = self.trainer.accelerator.validation_step(step_kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\accelerators\accelerator.py", line 236, in validation_step
    return self.training_type_plugin.validation_step(*step_kwargs.values())
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\pytorch_lightning\plugins\training_type\training_type_plugin.py", line 219, in validation_step
    return self.model.validation_step(*args, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\autograd\grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\models\diffusion\ddpm.py", line 368, in validation_step
    _, loss_dict_no_ema = self.shared_step(batch)
  File "C:\Dreambooth-SD-optimized\ldm\models\diffusion\ddpm.py", line 908, in shared_step
    loss = self(x, c)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\models\diffusion\ddpm.py", line 937, in forward
    c = self.get_learned_conditioning(c)
  File "C:\Dreambooth-SD-optimized\ldm\models\diffusion\ddpm.py", line 595, in get_learned_conditioning
    c = self.cond_stage_model.encode(c, embedding_manager=self.embedding_manager)
  File "C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py", line 324, in encode
    return self(text, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py", line 319, in forward
    z = self.transformer(input_ids=tokens, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py", line 297, in transformer_forward
    return self.text_model(
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py", line 258, in text_encoder_forward
    hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids, embedding_manager=embedding_manager)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py", line 180, in embedding_forward
    inputs_embeds = self.token_embedding(input_ids)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\modules\sparse.py", line 158, in forward
    return F.embedding(
  File "C:\Users\User\miniconda3\envs\ldm\lib\site-packages\torch\nn\functional.py", line 2044, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper__index_select)

修复建议 #

从报错栈可以看出，问题出在文本编码器的embedding阶段：input_ids在CPU，而token_embedding层的权重在CUDA设备上，导致张量设备不匹配。可以按以下步骤修复：

• 修改embedding_forward函数：
  打开C:\Dreambooth-SD-optimized\ldm\modules\encoders\modules.py，找到embedding_forward函数（对应报错行180），在调用self.token_embedding前，将input_ids移到模型权重所在设备：

  # 在inputs_embeds = self.token_embedding(input_ids)前添加
  input_ids = input_ids.to(self.token_embedding.weight.device)
  inputs_embeds = self.token_embedding(input_ids)
  

• 确保batch数据设备一致：
  打开C:\Dreambooth-SD-optimized\ldm\models\diffusion\ddpm.py，找到shared_step函数（对应报错行908），在开头添加代码将输入数据移到模型所在设备：

  x = x.to(self.device)
  c = c.to(self.device)
  

• 检查Trainer配置：
  确认训练器明确指定使用GPU加速，比如在main.py中：

  trainer = Trainer(accelerator='gpu', devices=1)
  

内容的提问来源于stack exchange，提问作者Zechlyr