Move training_output validation to after train_step_end

[kandluis]

What does this PR do?

See comments in https://github.com/PyTorchLightning/pytorch-lightning/pull/7779/files#r646750001.

We have a downstream LightningModule that does not conform to these expectations. It defines a trainign_step_end which returns the loss (after accumulating partial outputs from training_step) but where each training_step does not return a loss or a Mapping containing a key loss.

This still works with automatic optimization since the training_step_end function collects the results from training_step and itself returns a torch.Tensor corresponding to the loss.

So the question is: If I do define training_step_end, should I be allowed to not return a loss from training_step as long as I return it on training_step_end?

Assuming an affirmative to the above, this PR makes it so the validation occurs after training_step_end.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[ananthsub]

Thanks for the fix! I agree, we should run the check after training_step_end to allow for this

[codecov]

Codecov Report

Merging #7868 (09a73c1) into master (20f37b8) will decrease coverage by 6%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #7868    +/-   ##
=======================================
- Coverage      93%     87%    -6%     
=======================================
  Files         202     202            
  Lines       13123   13123            
=======================================
- Hits        12156   11402   -754     
- Misses        967    1721   +754     

[carmocca]

Do you want to add a test similar to this one for order?

https://github.com/PyTorchLightning/pytorch-lightning/blob/20f37b85b68f9903df8d61e79fcebdbadacf6422/tests/trainer/loops/test_evaluation_loop.py#L49-L71

[pep8speaks]

Hello @kandluis! Thanks for updating this PR.

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2021-06-08 00:06:29 UTC

[kandluis]

Do you want to add a test similar to this one for order?

https://github.com/PyTorchLightning/pytorch-lightning/blob/20f37b85b68f9903df8d61e79fcebdbadacf6422/tests/trainer/loops/test_evaluation_loop.py#L49-L71

Added a functional test for this behavior - let me know if you prefer I explicitly mock out the calls and track them as in the unit test you suggested.

[tchaton]

LGTM !