mlr3torch supports a number of PipeOp building blocks to define neural networks, including preprocessing and data augmentation steps.
PipeOp