Possible that the reduce function after validation_step might make it hard to save outputs such as videos, audios, etc... Need to evaluate whether it makes sense to remove.