DEV Community

CloudGPU
CloudGPU

Posted on

Why does my PyTorch DataLoader fail when a worker touches CUDA?

#pytorch #machinelearning #python #gpu

I work on CloudGPU. Keep dataset decoding and collation on the CPU, then move each batch to CUDA in the main training process. If a worker initializes CUDA after a fork, adding more workers can expose a process-initialization problem rather than a model problem.

1. Find the CUDA operation inside the worker

Inspect Dataset.__getitem__, transforms, collate_fn and worker_init_fn. Remove .cuda(), .to("cuda") and accelerator model construction from those paths. Also check whether a tensor was created on CUDA when the dataset object was built.

Use num_workers=0 temporarily to obtain a clearer traceback. It is a diagnostic step, not a recommendation about throughput.

2. Give the loader CPU data and an explicit process context

Save this as loader_boundary.py in an environment with a suitable PyTorch CUDA build:

import torch
from torch.utils.data import DataLoader, Dataset


class Samples(Dataset):
    def __len__(self):
        return 12

    def __getitem__(self, index):
        return torch.tensor([index], dtype=torch.float32)


def main():
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available in this environment")
    loader = DataLoader(
        Samples(), batch_size=3, num_workers=2,
        multiprocessing_context="spawn", pin_memory=True,
    )
    for cpu_batch in loader:
        assert cpu_batch.device.type == "cpu"
        gpu_batch = cpu_batch.to("cuda", non_blocking=True)
        assert gpu_batch.device.type == "cuda"


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode 
python3 loader_boundary.py
Enter fullscreen mode Exit fullscreen mode

Keep the dataset class at module scope and the execution behind the main guard. Run the example as a script, rather than assuming an interactive notebook can spawn the same definitions. CPU-only workers still need review for picklable objects and import-time side effects.

3. Restore your real data pipeline one part at a time

Add the real decoder first, then transforms, then collation. Keep the device transfer in the main loop. This isolates the first change that reintroduces CUDA into a worker. This answer does not diagnose every worker crash: unreadable inputs and CPU resource exhaustion need separate investigation.

The official DataLoader documentation recommends CPU batches and memory pinning instead of returning CUDA tensors from multiprocess loading. The multiprocessing guidance explains accelerator initialization and process start methods.

For rental-side execution boundaries, see the CloudGPU integration guide.

Related CloudGPU video: This Video Was Made by AI. Now Run AI Yourself. It provides remote-compute context, not a demonstration or validation of this loader example.

This is a documentation-based diagnostic recipe. The code was syntax-checked locally, not executed on a GPU; no training or performance result is claimed.

Top comments (0)