Google Colab GPU环境下DALLE-pytorch运行RuntimeError求助

问题场景 #

在Google Colab GPU环境运行DALLE-pytorch相关notebook时，执行以下代码：

with torch.no_grad():
    generated_images = vae.decode(generated_image_codes)

触发如下错误：

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)

解决方案 #

• 检查输入张量设备：确认generated_image_codes是否在CUDA设备上，执行以下代码查看并转换：

  print(generated_image_codes.device)  # 查看当前设备
  if generated_image_codes.device != 'cuda':
      generated_image_codes = generated_image_codes.to('cuda')
  

• 检查VAE模型设备：确保VAE模型本身运行在CUDA上，执行：

  print(vae.device)  # 查看模型当前设备
  if vae.device != 'cuda':
      vae = vae.to('cuda')
  

• 排查前置代码的设备转移逻辑：检查生成generated_image_codes的代码块，是否有将张量意外转移到CPU的操作（比如用.cpu()或者未指定设备的张量初始化），确保所有和模型交互的张量都统一在CUDA设备上。

内容的提问来源于stack exchange，提问作者Higo Felipe Silva Pires