Hm, I am not sure if that’s possible since the API is mainly targeted for DL training where the GPU is busy running the model (also, some people like to put transforms in there that are not implemented on GPUs, e.g., PIL or OpenCV stuff). I think in your case, it’s maybe not a good use case for the data loader with custom transform. I would use maybe the data loader just for iterating and do the processing then via a manual function like you did.