adapting `_validation_step`

Description of feature

There are two modifications which we need

1. we want to be able to pass kwargs to the solver as in predict. as we want to need them in every validatino iteration, we need to pass them once in prepare_validation_data, and then save them as attributes in the CellFlowTrainer
2. we want to pass the model to ComputationCallback, such that people can do predictions themselves, i.e. in addition to x_true, and x_predicted (in most cases this will just be absorped/ignored, e.g. in the ComputationCallback)