PinMemoryTransform¶
- class torchrl.envs.transforms.PinMemoryTransform[source]¶
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.
