Grain workers allocates GPU memory

When setting the GrainDataLoader worker_count > 1, each worker will use 500MB of GPU memory. When having many worker, this quickly depletes the amount of available memory, even if all processing is done on the CPU.

For example:

I have tried to avoid this by preventing CUDA to be used for sub processes like this:

```
import os
import multiprocessing
if multiprocessing.parent_process() is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["JAX_PLATFORMS"] = "cpu"
```

The root cause for this seems to be any reference to jax. When jax is imported, it seems to want to allocate some GPU memory. If we stick to numpy, this does not happen. We do need to use jax for resizing image/video tensors in a Transform operation.