PyTorch Lightning GPU模式CUDA初始化错误，CPU正常求排查

环境与基础排查 #

• 先验证CUDA环境是否正常，在代码开头添加以下代码确认：

import torch
print("CUDA可用:", torch.cuda.is_available())
print("GPU数量:", torch.cuda.device_count())
print("当前GPU:", torch.cuda.get_device_name(0))

若输出异常（比如CUDA可用: False），优先修复环境：Colab中可重启运行时，确保安装的PyTorch与PyTorch Lightning版本兼容（推荐PyTorch 2.x搭配Lightning 2.x）。

DataLoader参数调整 #

强化学习场景下的DataLoader常因参数设置不当导致GPU初始化失败，调整为以下配置：

from torch.utils.data import DataLoader

train_loader = DataLoader(
    your_dataset,
    batch_size=32,  # 根据显存调整合适大小
    shuffle=True,
    num_workers=2,  # Colab中建议设为2，避免资源过载
    pin_memory=True,  # 加速CPU到GPU的数据传输
    drop_last=True
)

同时确保Dataset的__getitem__方法中，所有张量默认保留在CPU，不要手动转CUDA，交给PyTorch Lightning统一管理设备转移。

Trainer配置优化 #

不要仅设置accelerator='gpu'，补充完整设备参数，避免自动多GPU适配出错：

from pytorch_lightning import Trainer

trainer = Trainer(
    accelerator="gpu",
    devices=1,  # 指定单GPU运行，新手阶段避免多GPU复杂度
    max_epochs=10,
    enable_progress_bar=True
)

同时删除代码中手动设置CUDA设备的语句（如torch.cuda.set_device()），让Lightning自动处理设备分配。

显存清理与冲突解决 #

Colab的GPU可能残留之前进程的显存占用，运行代码前先清空显存：

import torch
torch.cuda.empty_cache()

若问题依旧，重启Colab运行时（Runtime -> Restart runtime）后重新执行代码。

强化学习代码设备一致性检查 #

强化学习中的 replay buffer 等组件不要手动将数据移到GPU，保持在CPU即可，采样后由Lightning自动转移到GPU。例如避免以下错误写法：

# 错误：手动将缓冲区数据转CUDA
self.replay_buffer = [item.to("cuda") for item in collected_data]

改为让模型在forward时自动接收Lightning分配的设备张量。

简化代码定位问题 #

若以上方法无效，将代码简化为最小可复现版本：保留简单模型、基础DataLoader和Trainer配置，验证GPU是否能正常运行。若简化后正常，再逐步添加强化学习逻辑，定位报错的具体模块。

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