PinMemoryTransform

class torchrl.envs.transforms.PinMemoryTransform

Calls pin_memory on the tensordict to facilitate writing on CUDA devices.

forward(tensordict: TensorDictBase) → TensorDictBase

Reads the input tensordict, and for the selected keys, applies the transform.