PyTorch Dataset的__getitem__无法返回布尔变量的问题排查

1. generator指定CUDA设备的报错原因 #

PyTorch DataLoader在多进程模式（num_workers > 0）下，随机数生成器（generator）必须绑定CPU设备。这是因为多进程通过fork子进程实现，而CUDA上下文无法跨fork的子进程共享。如果强行将generator设为CUDA设备，PyTorch的内置校验会直接抛出Expected a 'cpu' device type for generator but found 'cuda'错误——生成器需要在主进程的CPU上生成随机数，再传递给子进程使用。

2. 返回Python布尔值触发“Cannot re-initialize CUDA”错误的原因 #

当Dataset的__getitem__返回Python原生布尔值时，DataLoader的多进程处理逻辑会触发CUDA初始化冲突：

• 多进程模式下，子进程会尝试重新初始化CUDA上下文，但Python原生对象（如bool）在跨进程传递时，会触发隐式的CUDA资源复制操作，导致子进程重复初始化CUDA，进而抛出Cannot re-initialize CUDA in forked subprocess错误。
• 而返回张量时，PyTorch有专门的跨进程张量传递机制（如共享内存、序列化处理），能避免这种重复初始化的问题。

3. 换成CPU上的torch.bool张量解决问题的原因 #

torch.bool是PyTorch原生张量类型，即使在CPU上，PyTorch也能通过内置的多进程数据传递机制正确处理：

• 张量会被正确序列化或通过共享内存传递给子进程，不会触发Python原生对象那种导致CUDA重复初始化的隐式操作。
• CPU张量不需要依赖CUDA上下文，子进程处理时不会涉及CUDA初始化的冲突。

4. 测试代码与原代码generator设备设置表现相反的原因 #

这种差异大概率源于两者的num_workers设置不同：

• 如果测试代码中num_workers=0（单进程模式），此时generator可以绑定CUDA设备——因为不需要fork子进程，主进程的CUDA上下文可以直接被使用，不会触发设备校验错误。
• 而原代码中num_workers>0（多进程模式），必须使用CPU的generator，所以设置CUDA会报错。进程模式的差异直接导致了generator设备设置的表现相反。

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