Amend docstring and add test for Flatten module (#42084)
Summary:
I've noticed when PR https://github.com/pytorch/pytorch/issues/22245 introduced `nn.Flatten`, the docstring had a bug where it wouldn't render properly on the web, and this PR addresses that. Additionally, it adds a unit test for this module.
**Actual**
**Expected**
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42084
Reviewed By: mrshenli
Differential Revision: D22756662
Pulled By: ngimel
fbshipit-source-id: 60c58c18c9a68854533196ed6b9e9fb0d4f83520
