运行Karlo扩散模型时，如何排查索引用CPU的RuntimeError？

• 追踪张量设备定位问题源
  在代码中所有涉及索引操作的位置（比如tensor[idx]这类语句）前后添加打印代码，输出索引和被索引张量的设备信息：

  print(f"Index tensor device: {idx.device}, Target tensor device: {tensor.device}")
  

  同时确认CUDA环境是否正常生效：

  import torch
  import os
  print(f"CUDA available: {torch.cuda.is_available()}")
  print(f"Visible GPU: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
  print(f"Current GPU: {torch.cuda.current_device()}")
  

  这样能精准定位到哪一行的张量设备不匹配。

• 排查隐性CPU张量生成

  • 检查从Numpy转换的张量：如果索引是由np.array()生成后转成PyTorch张量，务必指定设备，比如torch.tensor(numpy_index, device="cuda")，否则默认在CPU。
  • 检查数据加载流程：数据集预处理、DataLoader输出的张量是否都移到了GPU，比如DataLoader的pin_memory=True配合后续data.to("cuda")操作；PIL图像转张量时，要确保最终张量设备是GPU。
  • 检查模型内部固定张量：比如嵌入层的字典、掩码张量、时间步张量等初始化时的固定参数，这些可能默认在CPU，需手动调用.to("cuda")移到GPU。

• 检查预训练权重与第三方组件

  • 加载预训练权重时，明确指定map_location参数：

    torch.load(weights_path, map_location="cuda:0")
    

    避免权重默认加载到CPU，导致模型部分参数在CPU运行。
  • 排查模型依赖的第三方库组件：比如扩散模型的时间步生成、注意力机制的掩码，有些库内部生成的张量默认在CPU，需要手动移至GPU。

• 多GPU环境的设备同步问题

  • 确认指定的GPU（CUDA_VISIBLE_DEVICES=1）是否被正确识别，用torch.cuda.device_count()查看可用GPU数量，确保当前使用的是目标GPU。
  • 如果使用了DataParallel或DistributedDataParallel，需确保索引张量被同步到模型所在的GPU设备，避免跨GPU索引操作。

• 清理遗漏的CPU指定代码

  • 检查代码中的条件分支、循环逻辑：有些cpu()调用可能藏在分支里（比如调试用的代码、特定场景的处理逻辑），需要全局搜索并替换。
  • 检查张量创建函数：torch.arange()、torch.zeros()、torch.ones()等函数默认在CPU创建张量，必须显式添加device="cuda"参数。

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