Enable passing result from *_step Model Hook to corresponding *_batch_end Callback
🚀 Feature
Give users the option to pass a result at the end of *_step directly to the corresponding *_batch_end Callback.
Example:
validation_step outputs prediction and probabilities. Pass these to validation_batch_end user-defined Callback for advanced logging.
Motivation
This will remove the need for some boilerplate in the model code and also remove the need for calling the model's forward() on the batch inside the callback just to do logging.