CPU-side ColorJitter is a dataloader bottleneck — could the GPU path be restored?

DROID action SFT currently runs ColorJitter inside the dataloader workers, which is
computationally expensive.

1. In principle: ColorJitter is a random augmentation to begin with, so moving it to
   the GPU should only introduce CPU/CUDA kernel-level numerical differences, without
   changing the training semantics — is that right?

2. In practice: ColorJitter and the tail after it can be moved out of the dataloader
   worker and run in the main process on the GPU. The pipeline stage and the ordering stay exactly
   the same; only the execution device changes. The worker also passes along its RNG
   state, which keeps the final difference limited to CPU/CUDA kernel rounding.
   The cost is that worker-to-main-process traffic doubles, since the tensor returned
   by the worker goes from [T, C, 3H/2, W] to [3T, C, H, W].

Does this approach look viable?

Start by locating the DROID action SFT dataloader and its ColorJitter pipeline, then trace where the tail stage and worker RNG state are handled. Compare the proposed worker-to-main-process transfer with the current ordering; done means the GPU path is shown viable without semantic changes and its traffic or performance cost is measured.