Skip to content

Add pin_memory as a DataPipe #1013

@ejguan

Description

@ejguan

🚀 The feature

In previous DataLoader, we relies on the argument of pin_memory to launch a thread to move Tensor from CPU to GPU shared memory. This feature should be implemented as a DataPipe with is_replicable() -> False to keep it in the main process.

This should be easily achieved by doing the similar thing as prefetch with buffer size 1.

Motivation, pitch

Feature parity with DataLoader
This DataPipe can also become an indicator that the subsequent operations are on GPU.

Alternatives

No response

Additional context

No response

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions